"""
unyt_array class.
"""
import copy
import re
import sys
import warnings
from functools import lru_cache
from importlib.metadata import version
from numbers import Number as numeric_type
import numpy as np
from numpy import (
absolute,
add,
arccos,
arccosh,
arcsin,
arcsinh,
arctan,
arctan2,
arctanh,
bitwise_and,
bitwise_or,
bitwise_xor,
cbrt,
ceil,
clip,
conj,
copysign,
cos,
cosh,
deg2rad,
divide,
divmod as divmod_,
equal,
exp,
exp2,
expm1,
fabs,
floor,
floor_divide,
fmax,
fmin,
fmod,
frexp,
greater,
greater_equal,
heaviside,
hypot,
invert,
iscomplex,
isfinite,
isinf,
isnan,
isnat,
isreal,
ldexp,
left_shift,
less,
less_equal,
log,
log1p,
log2,
log10,
logaddexp,
logaddexp2,
logical_and,
logical_not,
logical_or,
logical_xor,
matmul,
maximum,
minimum,
mod,
modf,
multiply,
negative,
nextafter,
not_equal,
ones_like,
positive,
power,
rad2deg,
reciprocal,
remainder,
right_shift,
rint,
sign,
signbit,
sin,
sinh,
spacing,
sqrt,
square,
subtract,
tan,
tanh,
true_divide,
trunc,
)
from packaging.version import Version
from sympy import Rational
from unyt._on_demand_imports import _astropy, _dask, _pint
from unyt._pint_conversions import convert_pint_units
from unyt._unit_lookup_table import default_unit_symbol_lut
from unyt.dimensions import angle, temperature
from unyt.equivalencies import equivalence_registry
from unyt.exceptions import (
InvalidUnitEquivalence,
InvalidUnitOperation,
IterableUnitCoercionError,
MKSCGSConversionError,
SymbolNotFoundError,
UnitConversionError,
UnitOperationError,
UnitsNotReducible,
)
from unyt.unit_object import Unit, _check_em_conversion, _em_conversion
from unyt.unit_registry import (
UnitRegistry,
_correct_old_unit_registry,
_sanitize_unit_system,
default_unit_registry,
)
from unyt.unit_symbols import delta_degC, delta_degF
from ._deprecation import warn_deprecated
NULL_UNIT = Unit()
POWER_MAPPING = {multiply: lambda x: x, divide: lambda x: 2 - x}
DISALLOWED_DTYPES = (
"S", # bytestring
"a", # bytestring
"U", # (unicode) bytes
"O", # Python object
"M", # datetime
"m", # timedelta
)
__doctest_requires__ = {
("unyt_array.from_pint", "unyt_array.to_pint"): ["pint"],
("unyt_array.from_astropy", "unyt_array.to_astropy"): ["astropy"],
}
_COPY_IF_NEEDED = None if Version(version("numpy")) >= Version("2.0.0dev0") else False
# This is partially adapted from the following SO thread
# https://stackoverflow.com/questions/41668588/regex-to-match-scientific-notation
_NUMB_PATTERN = (
r"[+/-]?(?:((?:\d\.?\d*[Ee][+\-]?\d+)|(?:\d+\.\d*|\d*\.\d+))|\d+|inf\s|nan\s)"
)
# *all* greek letters are considered valid unit string elements.
# This may be an overshoot. We rely on unyt.Unit to do the actual validation
_UNIT_PATTERN = r"((\s*[*/]\s*)?[α-ωΑ-Ωa-zA-Z]+(\*\*([+-]?\d+|\([+-]?\d+\)))?)+"
_QUAN_PATTERN = rf"{_NUMB_PATTERN}\s*{_UNIT_PATTERN}"
_NUMB_REGEXP = re.compile(_NUMB_PATTERN)
_UNIT_REGEXP = re.compile(_UNIT_PATTERN)
_QUAN_REGEXP = re.compile(_QUAN_PATTERN)
def _iterable(obj):
try:
len(obj)
except Exception:
return False
return True
@lru_cache(maxsize=128, typed=False)
def _sqrt_unit(unit):
return 1, unit**0.5
@lru_cache(maxsize=128, typed=False)
def _cbrt_unit(unit):
return 1, unit ** (1.0 / 3.0)
@lru_cache(maxsize=128, typed=False)
def _multiply_units(unit1, unit2):
try:
ret = (unit1 * unit2).simplify()
except SymbolNotFoundError:
# Some operators are not natively commutative when operands are
# defined within different unit registries, and conversion
# is defined one way but not the other.
ret = (unit2 * unit1).simplify()
return ret.as_coeff_unit()
@lru_cache(maxsize=128, typed=False)
def _preserve_units(unit1, unit2=None):
if unit2 is None or unit1.dimensions is not temperature:
return 1, unit1
if unit1.base_offset == 0.0 and unit2.base_offset != 0.0:
return 1, unit2
return 1, unit1
@lru_cache(maxsize=128, typed=False)
def _difference_units(unit1, unit2=None):
if unit1.dimensions is not temperature:
return _preserve_units(unit1, unit2)
s1 = repr(unit1)
if unit2 is not None and unit2 != unit1:
s2 = repr(unit2)
if s1 in s2 and s2.startswith("delta_"):
return 1, unit1
elif s2 in s1 and s1.startswith("delta_"):
return 1, unit2
else:
raise InvalidUnitOperation(
"Quantities with units of Fahrenheit or Celsius "
"cannot be multiplied, divided, subtracted or "
"added with data that has different units."
)
if unit1.base_offset == 0.0:
return 1, unit1
if s1 == "degF":
return 1, delta_degF
elif s1 == "degC":
return 1, delta_degC
else:
# This is supposed to be unreachable
raise RuntimeError(
"Could not determine difference temperature units "
f"in operation ({unit1} - {unit2}).\n"
"If you see this error please file an issue at "
"https://github.com/yt-project/unyt/issues/new"
)
@lru_cache(maxsize=128, typed=False)
def _power_unit(unit, power):
return 1, unit**power
@lru_cache(maxsize=128, typed=False)
def _square_unit(unit):
return 1, unit * unit
@lru_cache(maxsize=128, typed=False)
def _divide_units(unit1, unit2):
try:
ret = (unit1 / unit2).simplify()
except SymbolNotFoundError:
ret = (1 / (unit2 / unit1).simplify()).units
return ret.as_coeff_unit()
@lru_cache(maxsize=128, typed=False)
def _reciprocal_unit(unit):
return 1, unit**-1
def _passthrough_unit(unit, unit2=None):
return 1, unit
def _return_without_unit(unit, unit2=None):
return 1, None
def _arctan2_unit(unit1, unit2):
return 1, NULL_UNIT
def _comparison_unit(unit1, unit2=None):
return 1, None
def _invert_units(unit):
raise TypeError("Bit-twiddling operators are not defined for unyt_array instances")
def _bitop_units(unit1, unit2):
raise TypeError("Bit-twiddling operators are not defined for unyt_array instances")
def _coerce_iterable_units(input_object, registry=None):
if isinstance(input_object, np.ndarray):
ret = input_object
elif _iterable(input_object):
if any(isinstance(o, unyt_array) for o in input_object):
ff = getattr(input_object[0], "units", NULL_UNIT)
if any(ff != getattr(_, "units", NULL_UNIT) for _ in input_object):
ret = []
for datum in input_object:
try:
ret.append(datum.in_units(ff.units))
except UnitConversionError:
raise IterableUnitCoercionError(str(input_object))
ret = unyt_array(np.array(ret), ff, registry=registry)
# This will create a copy of the data in the iterable.
else:
ret = unyt_array(np.array(input_object), ff, registry=registry)
else:
ret = np.asarray(input_object)
else:
ret = np.asarray(input_object)
if ret.dtype.char in DISALLOWED_DTYPES:
raise IterableUnitCoercionError(str(input_object))
return ret
def _sanitize_units_convert(possible_units, registry):
if isinstance(possible_units, Unit):
return possible_units
# let Unit() try to parse this if it's not already a Unit
unit = Unit(possible_units, registry=registry)
return unit
def _apply_power_mapping(ufunc, in_unit, in_size, in_shape, input_kwarg_dict):
# a reduction of a multiply or divide corresponds to
# a repeated product which we implement as an exponent
mul = 1
power_map = POWER_MAPPING[ufunc]
if input_kwarg_dict.get("axis", None) is not None:
unit = in_unit ** (power_map(in_shape[input_kwarg_dict["axis"]]))
else:
unit = in_unit ** (power_map(in_size))
return mul, unit
unary_operators = (
negative,
absolute,
rint,
sign,
conj,
exp,
exp2,
log,
log2,
log10,
expm1,
log1p,
sqrt,
cbrt,
square,
reciprocal,
sin,
cos,
tan,
arcsin,
arccos,
arctan,
sinh,
cosh,
tanh,
arcsinh,
arccosh,
arctanh,
deg2rad,
rad2deg,
invert,
logical_not,
isreal,
iscomplex,
isfinite,
isinf,
isnan,
signbit,
floor,
ceil,
trunc,
modf,
frexp,
fabs,
spacing,
positive,
isnat,
ones_like,
)
binary_operators = (
add,
subtract,
multiply,
divide,
logaddexp,
logaddexp2,
true_divide,
power,
remainder,
mod,
arctan2,
hypot,
bitwise_and,
bitwise_or,
bitwise_xor,
left_shift,
right_shift,
greater,
greater_equal,
less,
less_equal,
not_equal,
equal,
logical_and,
logical_or,
logical_xor,
maximum,
minimum,
fmax,
fmin,
copysign,
nextafter,
ldexp,
fmod,
divmod_,
heaviside,
)
trigonometric_operators = (sin, cos, tan)
multiple_output_operators = {modf: 2, frexp: 2, divmod_: 2}
LARGE_INPUT = {4: 16777217, 8: 9007199254740993}
class unyt_array(np.ndarray):
"""
An ndarray subclass that attaches a symbolic unit object to the array data.
Parameters
----------
input_array : iterable
A tuple, list, or array to attach units to
units : String unit name, unit symbol object, or astropy unit
The units of the array. Powers must be specified using python
syntax (cm**3, not cm^3).
registry : :class:`unyt.unit_registry.UnitRegistry`
The registry to create units from. If units is already associated
with a unit registry and this is specified, this will be used instead
of the registry associated with the unit object.
dtype : numpy dtype or dtype name
The dtype of the array data. Defaults to the dtype of the input data,
or, if none is found, uses np.float64
bypass_validation : boolean
If True, all input validation is skipped. Using this option may produce
corrupted, invalid units or array data, but can lead to significant
speedups in the input validation logic adds significant overhead. If
set, units *must* be a valid unit object. Defaults to False.
name : string
The name of the array. Defaults to None. This attribute does not propagate
through mathematical operations, but is preserved under indexing
and unit conversions.
Examples
--------
>>> from unyt import unyt_array
>>> a = unyt_array([1, 2, 3], 'cm')
>>> b = unyt_array([4, 5, 6], 'm')
>>> a + b
unyt_array([401., 502., 603.], 'cm')
>>> b + a
unyt_array([4.01, 5.02, 6.03], 'm')
NumPy ufuncs will pass through units where appropriate.
>>> from unyt import g, cm
>>> import numpy as np
>>> a = (np.arange(8) - 4)*g/cm**3
>>> np.abs(a)
unyt_array([4, 3, 2, 1, 0, 1, 2, 3], 'g/cm**3')
and strip them when it would be annoying to deal with them.
>>> np.log10(np.arange(8)+1)
array([0. , 0.30103 , 0.47712125, 0.60205999, 0.69897 ,
0.77815125, 0.84509804, 0.90308999])
"""
_ufunc_registry = {
add: _preserve_units,
subtract: _difference_units,
multiply: _multiply_units,
divide: _divide_units,
logaddexp: _return_without_unit,
logaddexp2: _return_without_unit,
true_divide: _divide_units,
floor_divide: _divide_units,
negative: _passthrough_unit,
power: _power_unit,
remainder: _preserve_units,
mod: _preserve_units,
fmod: _preserve_units,
absolute: _passthrough_unit,
fabs: _passthrough_unit,
rint: _return_without_unit,
sign: _return_without_unit,
conj: _passthrough_unit,
exp: _return_without_unit,
exp2: _return_without_unit,
log: _return_without_unit,
log2: _return_without_unit,
log10: _return_without_unit,
expm1: _return_without_unit,
log1p: _return_without_unit,
sqrt: _sqrt_unit,
cbrt: _cbrt_unit,
square: _square_unit,
reciprocal: _reciprocal_unit,
sin: _return_without_unit,
cos: _return_without_unit,
tan: _return_without_unit,
sinh: _return_without_unit,
cosh: _return_without_unit,
tanh: _return_without_unit,
arcsin: _return_without_unit,
arccos: _return_without_unit,
arctan: _return_without_unit,
arctan2: _arctan2_unit,
arcsinh: _return_without_unit,
arccosh: _return_without_unit,
arctanh: _return_without_unit,
hypot: _preserve_units,
deg2rad: _return_without_unit,
rad2deg: _return_without_unit,
bitwise_and: _bitop_units,
bitwise_or: _bitop_units,
bitwise_xor: _bitop_units,
invert: _invert_units,
left_shift: _bitop_units,
right_shift: _bitop_units,
greater: _comparison_unit,
greater_equal: _comparison_unit,
less: _comparison_unit,
less_equal: _comparison_unit,
not_equal: _comparison_unit,
equal: _comparison_unit,
logical_and: _comparison_unit,
logical_or: _comparison_unit,
logical_xor: _comparison_unit,
logical_not: _return_without_unit,
maximum: _preserve_units,
minimum: _preserve_units,
fmax: _preserve_units,
fmin: _preserve_units,
isreal: _return_without_unit,
iscomplex: _return_without_unit,
isfinite: _return_without_unit,
isinf: _return_without_unit,
isnan: _return_without_unit,
signbit: _return_without_unit,
copysign: _passthrough_unit,
nextafter: _preserve_units,
modf: _passthrough_unit,
ldexp: _bitop_units,
frexp: _return_without_unit,
floor: _passthrough_unit,
ceil: _passthrough_unit,
trunc: _passthrough_unit,
spacing: _passthrough_unit,
positive: _passthrough_unit,
divmod_: _passthrough_unit,
isnat: _return_without_unit,
heaviside: _preserve_units,
matmul: _multiply_units,
clip: _passthrough_unit,
}
__array_priority__ = 2.0
def __new__(
cls,
input_array,
units=None,
registry=None,
dtype=None,
*,
bypass_validation=False,
name=None,
):
input_units = units
if bypass_validation is True:
if dtype is None:
dtype = input_array.dtype
obj = input_array.view(type=cls, dtype=dtype)
obj.units = input_units
if registry is not None:
obj.units.registry = registry
obj.name = name
return obj
if isinstance(input_array, unyt_array):
ret = input_array.view(cls)
if input_units is None:
if registry is None:
ret.units = input_array.units
else:
units = Unit(str(input_array.units), registry=registry)
ret.units = units
elif isinstance(input_units, Unit):
ret.units = input_units
else:
ret.units = Unit(input_units, registry=registry)
ret.name = name
return ret
elif isinstance(input_array, np.ndarray):
pass
elif _iterable(input_array) and input_array:
if isinstance(input_array[0], unyt_array):
return _coerce_iterable_units(input_array, registry)
# Input array is an already formed ndarray instance
# We first cast to be our class type
obj = np.asarray(input_array, dtype=dtype).view(cls)
# Check units type
if input_units is None:
# Nothing provided. Make dimensionless...
units = Unit()
elif isinstance(input_units, Unit):
if registry and registry is not input_units.registry:
units = Unit(str(input_units), registry=registry)
else:
units = input_units
else:
# units kwarg set, but it's not a Unit object.
# don't handle all the cases here, let the Unit class handle if
# it's a str.
units = Unit(input_units, registry=registry)
# Attach the units and name
obj.units = units
obj.name = name
return obj
def __repr__(self):
return np.array_repr(self)
def __str__(self):
return str(self.view(np.ndarray)) + " " + str(self.units)
def __format__(self, format_spec):
return f"{self.d.__format__(format_spec)} {self.units}"
#
# Start unit conversion methods
#
def convert_to_units(self, units, equivalence=None, **kwargs):
"""
Convert the array to the given units in-place.
Optionally, an equivalence can be specified to convert to an
equivalent quantity which is not in the same dimensions.
Parameters
----------
units : Unit object or string
The units you want to convert to.
equivalence : string, optional
The equivalence you wish to use. To see which equivalencies
are supported for this object, try the ``list_equivalencies``
method. Default: None
kwargs: optional
Any additional keyword arguments are supplied to the equivalence
Raises
------
If the provided unit does not have the same dimensions as the array
this will raise a UnitConversionError
Examples
--------
>>> from unyt import cm, km
>>> length = [3000, 2000, 1000]*cm
>>> length.convert_to_units('m')
>>> print(length)
[30. 20. 10.] m
"""
units = _sanitize_units_convert(units, self.units.registry)
if equivalence is None:
conv_data = _check_em_conversion(
self.units, units, registry=self.units.registry
)
if any(conv_data):
new_units, (conv_factor, offset) = _em_conversion(
self.units, conv_data, units
)
else:
new_units = units
(conv_factor, offset) = self.units.get_conversion_factor(
new_units, self.dtype
)
self.units = new_units
values = self.d
# if our dtype is an integer do the following somewhat awkward
# dance to change the dtype in-place. We can't use astype
# directly because that will create a copy and not update self
if self.dtype.kind in ("u", "i"):
# create a copy of the original data in floating point
# form, it's possible this may lose precision for very
# large integers
dsize = values.dtype.itemsize
if dsize == 1:
raise ValueError(
"Can't convert memory buffer in place. "
f"Input dtype ({self.dtype}) has a smaller itemsize than the "
"smallest floating point representation possible."
)
new_dtype = "f" + str(dsize)
large = LARGE_INPUT.get(dsize, 0)
if large and np.any(np.abs(values) > large):
warnings.warn(
f"Overflow encountered while converting to units '{new_units}'",
RuntimeWarning,
stacklevel=2,
)
float_values = values.astype(new_dtype)
# change the dtypes in-place, this does not change the
# underlying memory buffer
values.dtype = new_dtype
self.dtype = new_dtype
# actually fill in the new float values now that our
# dtype is correct
np.copyto(values, float_values)
values *= conv_factor
if offset:
np.subtract(values, offset, values)
else:
self.convert_to_equivalent(units, equivalence, **kwargs)
def convert_to_base(self, unit_system=None, equivalence=None, **kwargs):
"""
Convert the array in-place to the equivalent base units in
the specified unit system.
Optionally, an equivalence can be specified to convert to an
equivalent quantity which is not in the same dimensions.
Parameters
----------
unit_system : string, optional
The unit system to be used in the conversion. If not specified,
the configured base units are used (defaults to MKS).
equivalence : string, optional
The equivalence you wish to use. To see which equivalencies
are supported for this object, try the ``list_equivalencies``
method. Default: None
kwargs: optional
Any additional keyword arguments are supplied to the equivalence
Raises
------
If the provided unit does not have the same dimensions as the array
this will raise a UnitConversionError
Examples
--------
>>> from unyt import erg, s
>>> E = 2.5*erg/s
>>> E.convert_to_base("mks")
>>> E
unyt_quantity(2.5e-07, 'W')
"""
self.convert_to_units(
self.units.get_base_equivalent(unit_system),
equivalence=equivalence,
**kwargs,
)
def convert_to_cgs(self, equivalence=None, **kwargs):
"""
Convert the array and in-place to the equivalent cgs units.
Optionally, an equivalence can be specified to convert to an
equivalent quantity which is not in the same dimensions.
Parameters
----------
equivalence : string, optional
The equivalence you wish to use. To see which equivalencies
are supported for this object, try the ``list_equivalencies``
method. Default: None
kwargs: optional
Any additional keyword arguments are supplied to the equivalence
Raises
------
If the provided unit does not have the same dimensions as the array
this will raise a UnitConversionError
Examples
--------
>>> from unyt import Newton
>>> data = [1., 2., 3.]*Newton
>>> data.convert_to_cgs()
>>> data
unyt_array([100000., 200000., 300000.], 'dyn')
"""
self.convert_to_units(
self.units.get_cgs_equivalent(), equivalence=equivalence, **kwargs
)
def convert_to_mks(self, equivalence=None, **kwargs):
"""
Convert the array and units to the equivalent mks units.
Optionally, an equivalence can be specified to convert to an
equivalent quantity which is not in the same dimensions.
Parameters
----------
equivalence : string, optional
The equivalence you wish to use. To see which equivalencies
are supported for this object, try the ``list_equivalencies``
method. Default: None
kwargs: optional
Any additional keyword arguments are supplied to the equivalence
Raises
------
If the provided unit does not have the same dimensions as the array
this will raise a UnitConversionError
Examples
--------
>>> from unyt import dyne, erg
>>> data = [1., 2., 3.]*erg
>>> data
unyt_array([1., 2., 3.], 'erg')
>>> data.convert_to_mks()
>>> data
unyt_array([1.e-07, 2.e-07, 3.e-07], 'J')
"""
self.convert_to_units(self.units.get_mks_equivalent(), equivalence, **kwargs)
def in_units(self, units, equivalence=None, **kwargs):
"""
Creates a copy of this array with the data converted to the
supplied units, and returns it.
Optionally, an equivalence can be specified to convert to an
equivalent quantity which is not in the same dimensions.
Parameters
----------
units : Unit object or string
The units you want to get a new quantity in.
equivalence : string, optional
The equivalence you wish to use. To see which equivalencies
are supported for this object, try the ``list_equivalencies``
method. Default: None
kwargs: optional
Any additional keyword arguments are supplied to the
equivalence
Raises
------
If the provided unit does not have the same dimensions as the array
this will raise a UnitConversionError
Examples
--------
>>> from unyt import c, gram
>>> m = 10*gram
>>> E = m*c**2
>>> print(E.in_units('erg'))
8.987551787368176e+21 erg
>>> print(E.in_units('J'))
898755178736817.6 J
"""
units = _sanitize_units_convert(units, self.units.registry)
if equivalence is None:
conv_data = _check_em_conversion(
self.units, units, registry=self.units.registry
)
if any(conv_data):
new_units, (conversion_factor, offset) = _em_conversion(
self.units, conv_data, units
)
offset = 0
else:
new_units = units
(conversion_factor, offset) = self.units.get_conversion_factor(
new_units, self.dtype
)
dsize = max(2, self.dtype.itemsize)
if self.dtype.kind in ("u", "i"):
large = LARGE_INPUT.get(dsize, 0)
if large and np.any(np.abs(self.d) > large):
warnings.warn(
f"Overflow encountered while converting to units '{new_units}'",
RuntimeWarning,
stacklevel=2,
)
new_dtypekind = "c" if self.dtype.kind == "c" else "f"
new_dtype = np.dtype(new_dtypekind + str(dsize))
ret = np.asarray(self.ndview * conversion_factor, dtype=new_dtype)
if offset:
np.subtract(ret, offset, ret)
try:
new_array = type(self)(
ret, new_units, bypass_validation=True, name=self.name
)
except TypeError:
# subclasses might not take name as a kwarg
new_array = type(self)(ret, new_units, bypass_validation=True)
return new_array
else:
return self.to_equivalent(units, equivalence, **kwargs)
def to(self, units, equivalence=None, **kwargs):
"""
Creates a copy of this array with the data converted to the
supplied units, and returns it.
Optionally, an equivalence can be specified to convert to an
equivalent quantity which is not in the same dimensions.
.. note::
All additional keyword arguments are passed to the
equivalency, which should be used if that particular
equivalency requires them.
Parameters
----------
units : Unit object or string
The units you want to get a new quantity in.
equivalence : string, optional
The equivalence you wish to use. To see which
equivalencies are supported for this unitful
quantity, try the :meth:`list_equivalencies`
method. Default: None
kwargs: optional
Any additional keywoard arguments are supplied to the
equivalence
Raises
------
If the provided unit does not have the same dimensions as the array
this will raise a UnitConversionError
Examples
--------
>>> from unyt import c, gram
>>> m = 10*gram
>>> E = m*c**2
>>> print(E.to('erg'))
8.987551787368176e+21 erg
>>> print(E.to('J'))
898755178736817.6 J
"""
return self.in_units(units, equivalence=equivalence, **kwargs)
def to_value(self, units=None, equivalence=None, **kwargs):
"""
Creates a copy of this array with the data in the supplied
units, and returns it without units. Output is therefore a
bare NumPy array.
Optionally, an equivalence can be specified to convert to an
equivalent quantity which is not in the same dimensions.
.. note::
All additional keyword arguments are passed to the
equivalency, which should be used if that particular
equivalency requires them.
Parameters
----------
units : Unit object or string, optional
The units you want to get the bare quantity in. If not
specified, the value will be returned in the current units.
equivalence : string, optional
The equivalence you wish to use. To see which
equivalencies are supported for this unitful
quantity, try the :meth:`list_equivalencies`
method. Default: None
Examples
--------
>>> from unyt import km
>>> a = [3, 4, 5]*km
>>> print(a.to_value('cm'))
[300000. 400000. 500000.]
"""
if units is None:
v = self.value
else:
v = self.in_units(units, equivalence=equivalence, **kwargs).value
if isinstance(self, unyt_quantity):
return float(v)
else:
return v
def in_base(self, unit_system=None):
"""
Creates a copy of this array with the data in the specified unit
system, and returns it in that system's base units.
Parameters
----------
unit_system : string, optional
The unit system to be used in the conversion. If not specified,
the configured default base units of are used (defaults to MKS).
Examples
--------
>>> from unyt import erg, s
>>> E = 2.5*erg/s
>>> print(E.in_base("mks"))
2.5e-07 W
"""
us = _sanitize_unit_system(unit_system, self)
try:
conv_data = _check_em_conversion(
self.units, unit_system=us, registry=self.units.registry
)
except MKSCGSConversionError:
raise UnitsNotReducible(self.units, us)
if any(conv_data):
um = us.units_map
u = self.units
if u.dimensions in um and u.expr == um[self.units.dimensions]:
return self.copy()
to_units, (conv, offset) = _em_conversion(u, conv_data, unit_system=us)
else:
to_units = self.units.get_base_equivalent(unit_system)
conv, offset = self.units.get_conversion_factor(to_units, self.dtype)
ret = self.v * conv
if offset:
ret = ret - offset
return type(self)(ret, to_units)
def in_cgs(self):
"""
Creates a copy of this array with the data in the equivalent cgs units,
and returns it.
Returns
-------
unyt_array object with data in this array converted to cgs units.
Example
-------
>>> from unyt import Newton, km
>>> print((10*Newton/km).in_cgs())
10.0 g/s**2
"""
return self.in_base("cgs")
def in_mks(self):
"""
Creates a copy of this array with the data in the equivalent mks units,
and returns it.
Returns
-------
unyt_array object with data in this array converted to mks units.
Example
-------
>>> from unyt import mile
>>> print((1.*mile).in_mks())
1609.344 m
"""
return self.in_base("mks")
def convert_to_equivalent(self, unit, equivalence, **kwargs):
"""
Convert the array in-place to the specified units, assuming
the given equivalency. The dimensions of the specified units and the
dimensions of the original array need not match so long as there is an
appropriate conversion in the specified equivalency.
Parameters
----------
unit : string
The unit that you wish to convert to.
equivalence : string
The equivalence you wish to use. To see which equivalencies are
supported for this unitful quantity, try the
:meth:`list_equivalencies` method.
Examples
--------
>>> from unyt import K
>>> a = [10, 20, 30]*(1e7*K)
>>> a.convert_to_equivalent("keV", "thermal")
>>> a
unyt_array([ 8.6173324, 17.2346648, 25.8519972], 'keV')
"""
conv_unit = Unit(unit, registry=self.units.registry)
if self.units.same_dimensions_as(conv_unit):
self.convert_to_units(conv_unit)
return
this_equiv = equivalence_registry[equivalence](in_place=True)
if self.has_equivalent(equivalence):
this_equiv.convert(self, conv_unit.dimensions, **kwargs)
self.convert_to_units(conv_unit)
# set name to None since the semantic meaning has changed
self.name = None
else:
raise InvalidUnitEquivalence(equivalence, self.units, conv_unit)
def to_equivalent(self, unit, equivalence, **kwargs):
"""
Return a copy of the unyt_array in the units specified units, assuming
the given equivalency. The dimensions of the specified units and the
dimensions of the original array need not match so long as there is an
appropriate conversion in the specified equivalency.
Parameters
----------
unit : string
The unit that you wish to convert to.
equivalence : string
The equivalence you wish to use. To see which equivalencies are
supported for this unitful quantity, try the
:meth:`list_equivalencies` method.
Examples
--------
>>> from unyt import K
>>> a = 1.0e7*K
>>> print(a.to_equivalent("keV", "thermal"))
0.8617332401096504 keV
"""
conv_unit = Unit(unit, registry=self.units.registry)
if self.units.same_dimensions_as(conv_unit):
return self.in_units(conv_unit)
this_equiv = equivalence_registry[equivalence]()
if self.has_equivalent(equivalence):
new_arr = this_equiv.convert(self, conv_unit.dimensions, **kwargs)
return new_arr.in_units(conv_unit)
else:
raise InvalidUnitEquivalence(equivalence, self.units, unit)
def list_equivalencies(self):
"""
Lists the possible equivalencies associated with this unyt_array or
unyt_quantity.
Example
-------
>>> from unyt import km
>>> (1.0*km).list_equivalencies()
spectral: length <-> spatial_frequency <-> frequency <-> energy
schwarzschild: mass <-> length
compton: mass <-> length
"""
self.units.list_equivalencies()
def has_equivalent(self, equivalence):
"""
Check to see if this unyt_array or unyt_quantity has an equivalent
unit in *equiv*.
Example
-------
>>> from unyt import km, keV
>>> (1.0*km).has_equivalent('spectral')
True
>>> print((1*km).to_equivalent('MHz', equivalence='spectral'))
0.299792458 MHz
>>> print((1*keV).to_equivalent('angstrom', equivalence='spectral'))
12.39841931521966 Å
"""
return self.units.has_equivalent(equivalence)
def ndarray_view(self):
"""
Returns a view into the array as a numpy array
Returns
-------
View of this array's data.
Example
-------
>>> from unyt import km
>>> a = [3, 4, 5]*km
>>> a
unyt_array([3, 4, 5], 'km')
>>> a.ndarray_view()
array([3, 4, 5])
This function returns a view that shares the same underlying memory
as the original array.
>>> b = a.ndarray_view()
>>> b.base is a.base
True
>>> b[2] = 4
>>> b
array([3, 4, 4])
>>> a
unyt_array([3, 4, 4], 'km')
"""
return self.view(np.ndarray)
def to_ndarray(self):
"""
Creates a copy of this array with the unit information stripped
Example
-------
>>> from unyt import km
>>> a = [3, 4, 5]*km
>>> a
unyt_array([3, 4, 5], 'km')
>>> b = a.to_ndarray()
>>> b
array([3, 4, 5])
The returned array will contain a copy of the data contained in
the original array.
>>> a.base is not b.base
True
"""
return np.array(self)
def argsort(self, axis=-1, kind="quicksort", order=None):
"""
Returns the indices that would sort the array.
See the documentation of ndarray.argsort for details about the keyword
arguments.
Example
-------
>>> from unyt import km
>>> data = [3, 8, 7]*km
>>> print(np.argsort(data))
[0 2 1] km
>>> print(data.argsort())
[0 2 1]
"""
return self.view(np.ndarray).argsort(axis, kind, order)
@classmethod
def from_astropy(cls, arr, unit_registry=None):
"""
Convert an AstroPy "Quantity" to a unyt_array or unyt_quantity.
Parameters
----------
arr : AstroPy Quantity
The Quantity to convert from.
unit_registry : unyt.UnitRegistry, optional
A unyt unit registry to use in the conversion. If one is not
supplied, the default one will be used.
Example
-------
>>> from astropy.units import km
>>> unyt_quantity.from_astropy(km)
unyt_quantity(1., 'km')
>>> a = [1, 2, 3]*km
>>> a
<Quantity [1., 2., 3.] km>
>>> unyt_array.from_astropy(a)
unyt_array([1., 2., 3.], 'km')
"""
# Converting from AstroPy Quantity
try:
u = arr.unit
_arr = arr
except AttributeError:
u = arr
_arr = 1.0 * u
ap_units = []
for base, exponent in zip(u.bases, u.powers):
unit_str = base.to_string()
# we have to do this because AstroPy is silly and defines
# hour as "h"
if unit_str == "h":
unit_str = "hr"
ap_units.append(f"{unit_str}**({Rational(exponent)})")
ap_units = "*".join(ap_units)
if isinstance(_arr.value, np.ndarray) and _arr.shape != ():
return unyt_array(_arr.value, ap_units, registry=unit_registry)
else:
return unyt_quantity(_arr.value, ap_units, registry=unit_registry)
def to_astropy(self, **kwargs):
"""
Creates a new AstroPy quantity with the same unit information.
Example
-------
>>> from unyt import g, cm
>>> data = [3, 4, 5]*g/cm**3
>>> data.to_astropy()
<Quantity [3., 4., 5.] g / cm3>
"""
if self.units.is_dimensionless:
s_units = ""
else:
s_units = str(self.units)
return self.value * _astropy.units.Unit(s_units, **kwargs)
@classmethod
def from_pint(cls, arr, unit_registry=None):
"""
Convert a Pint "Quantity" to a unyt_array or unyt_quantity.
Parameters
----------
arr : Pint Quantity
The Quantity to convert from.
unit_registry : unyt.UnitRegistry, optional
A unyt unit registry to use in the conversion. If one is not
supplied, the default one will be used.
Examples
--------
>>> from pint import UnitRegistry
>>> import numpy as np
>>> ureg = UnitRegistry()
>>> a = np.arange(4)
>>> b = ureg.Quantity(a, "erg/cm**3")
>>> b
<Quantity([0 1 2 3], 'erg / centimeter ** 3')>
>>> c = unyt_array.from_pint(b)
>>> c
unyt_array([0, 1, 2, 3], 'erg/cm**3')
"""
p_units = []
for base, exponent in arr._units.items():
bs = convert_pint_units(base)
p_units.append(f"{bs}**({Rational(exponent)})")
p_units = "*".join(p_units)
if isinstance(arr.magnitude, np.ndarray):
return unyt_array(arr.magnitude, p_units, registry=unit_registry)
else:
return unyt_quantity(arr.magnitude, p_units, registry=unit_registry)
def to_pint(self, unit_registry=None):
"""
Convert a unyt_array or unyt_quantity to a Pint Quantity.
Parameters
----------
arr : unyt_array or unyt_quantity
The unitful quantity to convert from.
unit_registry : Pint UnitRegistry, optional
The Pint UnitRegistry to use in the conversion. If one is not
supplied, the default one will be used. NOTE: This is not
the same as a unyt.UnitRegistry object.
Examples
--------
>>> from unyt import cm, s
>>> a = 4*cm**2/s
>>> print(a)
4 cm**2/s
>>> a.to_pint()
<Quantity(4, 'centimeter ** 2 / second')>
"""
if unit_registry is None:
unit_registry = _pint.UnitRegistry()
powers_dict = self.units.expr.as_powers_dict()
units = []
for unit, pow in powers_dict.items():
# we have to do this because Pint doesn't recognize
# "yr" as "year"
if str(unit).endswith("yr") and len(str(unit)) in [2, 3]:
unit = str(unit).replace("yr", "year")
units.append(f"{unit}**({Rational(pow)})")
units = "*".join(units)
return unit_registry.Quantity(self.value, units)
@staticmethod
def from_string(s, unit_registry=None):
"""
Parse a string to a unyt_quantity object.
Parameters
----------
s: str
A string representing a single quantity.
unit_registry: unyt.UnitRegistry, optional
A unyt unit registry to use in the conversion. If one is not
supplied, the default one will be used.
Examples
--------
>>> from unyt import unyt_quantity
>>> unyt_quantity.from_string("1cm")
unyt_quantity(1, 'cm')
>>> unyt_quantity.from_string("+1e3 Mearth")
unyt_quantity(1000., 'Mearth')
>>> unyt_quantity.from_string("-10. kg")
unyt_quantity(-10., 'kg')
>>> unyt_quantity.from_string(".66\tum")
unyt_quantity(0.66, 'μm')
>>> unyt_quantity.from_string("42")
unyt_quantity(42, '(dimensionless)')
>>> unyt_quantity.from_string("1.0 g/cm**3")
unyt_quantity(1., 'g/cm**3')
"""
v = s.strip()
if re.fullmatch(_NUMB_REGEXP, v):
num = re.match(_NUMB_REGEXP, v).group()
unit = Unit()
elif re.fullmatch(_UNIT_REGEXP, v):
num = 1
unit = Unit(re.match(_UNIT_REGEXP, v).group())
elif not re.fullmatch(_QUAN_REGEXP, v):
raise ValueError(f"Received invalid quantity expression '{s}'.")
else:
res = re.search(_NUMB_REGEXP, v)
num = res.group()
res = re.search(_UNIT_REGEXP, v[res.span()[1] :])
unit = res.group().strip()
if unit.startswith(("/", "*")):
unit = f"1{unit}"
try:
num = int(num)
except ValueError:
num = float(num)
return num * Unit(unit, registry=unit_registry)
def to_string(self):
# this is implemented purely for symmetry's sake
return str(self)
#
# End unit conversion methods
#
def write_hdf5(self, filename, dataset_name=None, info=None, group_name=None):
r"""Writes a unyt_array to hdf5 file.
Parameters
----------
filename: string
The filename to create and write a dataset to
dataset_name: string
The name of the dataset to create in the file.
info: dictionary
A dictionary of supplementary info to write to append as attributes
to the dataset.
group_name: string
An optional group to write the arrays to. If not specified, the
arrays are datasets at the top level by default.
Examples
--------
>>> from unyt import cm
>>> a = [1,2,3]*cm
>>> myinfo = {'field':'dinosaurs', 'type':'field_data'}
>>> a.write_hdf5('test_array_data.h5', dataset_name='dinosaurs',
... info=myinfo) # doctest: +SKIP
"""
import pickle
from unyt._on_demand_imports import _h5py as h5py
if info is None:
info = {}
info["units"] = str(self.units)
lut = {}
for k, v in self.units.registry.lut.items():
if k not in default_unit_registry.lut:
lut[k] = v
info["unit_registry"] = np.void(pickle.dumps(lut))
if dataset_name is None:
dataset_name = "array_data"
f = h5py.File(filename, "a")
if group_name is not None:
if group_name in f:
g = f[group_name]
else:
g = f.create_group(group_name)
else:
g = f
if dataset_name in g.keys():
d = g[dataset_name]
# Overwrite without deleting if we can get away with it.
if d.shape == self.shape and d.dtype == self.dtype:
d[...] = self
for k in d.attrs.keys():
del d.attrs[k]
else:
del f[dataset_name]
d = g.create_dataset(dataset_name, data=self)
else:
d = g.create_dataset(dataset_name, data=self)
for k, v in info.items():
d.attrs[k] = v
f.close()
@classmethod
def from_hdf5(cls, filename, dataset_name=None, group_name=None):
r"""Attempts read in and convert a dataset in an hdf5 file into a
unyt_array.
Parameters
----------
filename: string
The filename to of the hdf5 file.
dataset_name: string
The name of the dataset to read from. If the dataset has a units
attribute, attempt to infer units as well.
group_name: string
An optional group to read the arrays from. If not specified, the
arrays are datasets at the top level by default.
"""
import pickle
from unyt._on_demand_imports import _h5py as h5py
if dataset_name is None:
dataset_name = "array_data"
f = h5py.File(filename, "r")
if group_name is not None:
g = f[group_name]
else:
g = f
dataset = g[dataset_name]
data = dataset[:]
units = dataset.attrs.get("units", "")
unit_lut = default_unit_symbol_lut.copy()
unit_lut_load = pickle.loads(dataset.attrs["unit_registry"].tobytes())
unit_lut.update(unit_lut_load)
f.close()
registry = UnitRegistry(lut=unit_lut, add_default_symbols=False)
return cls(data, units, registry=registry)
#
# Start convenience methods
#
@property
def value(self):
"""
Creates a copy of this array with the unit information stripped
Example
-------
>>> from unyt import km
>>> a = [3, 4, 5]*km
>>> a
unyt_array([3, 4, 5], 'km')
>>> b = a.value
>>> b
array([3, 4, 5])
The returned array will contain a copy of the data contained in
the original array.
>>> a.base is not b.base
True
"""
return np.array(self)
@property
def v(self):
"""
Creates a copy of this array with the unit information stripped
Example
-------
>>> from unyt import km
>>> a = [3, 4, 5]*km
>>> a
unyt_array([3, 4, 5], 'km')
>>> b = a.v
>>> b
array([3, 4, 5])
The returned array will contain a copy of the data contained in
the original array.
>>> a.base is not b.base
True
"""
return np.array(self)
@property
def ndview(self):
"""
Returns a view into the array as a numpy array
Returns
-------
View of this array's data.
Example
-------
>>> from unyt import km
>>> a = [3, 4, 5]*km
>>> a
unyt_array([3, 4, 5], 'km')
>>> a.ndview
array([3, 4, 5])
This function returns a view that shares the same underlying memory
as the original array.
>>> b = a.ndview
>>> b.base is a.base
True
>>> b[2] = 4
>>> b
array([3, 4, 4])
>>> a
unyt_array([3, 4, 4], 'km')
"""
return self.view(np.ndarray)
@property
def d(self):
"""
Returns a view into the array as a numpy array
Returns
-------
View of this array's data.
Example
-------
>>> from unyt import km
>>> a = [3, 4, 5]*km
>>> a
unyt_array([3, 4, 5], 'km')
>>> a.d
array([3, 4, 5])
This function returns a view that shares the same underlying memory
as the original array.
>>> b = a.d
>>> b.base is a.base
True
>>> b[2] = 4
>>> b
array([3, 4, 4])
>>> a
unyt_array([3, 4, 4], 'km')
"""
return self.view(np.ndarray)
@property
def unit_quantity(self):
"""
Return a quantity with a value of 1 and the same units as this array
Example
-------
>>> from unyt import km
>>> a = [4, 5, 6]*km
>>> a.unit_quantity
unyt_quantity(1, 'km')
>>> print(a + 7*a.unit_quantity)
[11 12 13] km
"""
return unyt_quantity(1, self.units)
@property
def uq(self):
"""
Return a quantity with a value of 1 and the same units as this array
Example
-------
>>> from unyt import km
>>> a = [4, 5, 6]*km
>>> a.uq
unyt_quantity(1, 'km')
>>> print(a + 7*a.uq)
[11 12 13] km
"""
return unyt_quantity(1, self.units)
@property
def unit_array(self):
"""
Return an array filled with ones with the same units as this array
Example
-------
>>> from unyt import km
>>> a = [4, 5, 6]*km
>>> a.unit_array
unyt_array([1, 1, 1], 'km')
>>> print(a + 7*a.unit_array)
[11 12 13] km
"""
return np.ones_like(self)
@property
def ua(self):
"""
Return an array filled with ones with the same units as this array
Example
-------
>>> from unyt import km
>>> a = [4, 5, 6]*km
>>> a.unit_array
unyt_array([1, 1, 1], 'km')
>>> print(a + 7*a.unit_array)
[11 12 13] km
"""
return np.ones_like(self)
def __getitem__(self, item):
ret = super().__getitem__(item)
if getattr(ret, "shape", None) == ():
ret = unyt_quantity(ret, bypass_validation=True, name=self.name)
ret.units = self.units
return ret
def __setitem__(self, item, value):
if hasattr(value, "units"):
if value.units != self.units and value.units != NULL_UNIT:
value = value.to(self.units)
super().__setitem__(item, value)
def __pow__(self, p, mod=None, /):
"""
Power function
"""
# see https://github.com/yt-project/unyt/issues/203
if p == 0.0:
ret = self.ua
ret.units = Unit("dimensionless")
return ret
else:
return super().__pow__(p, mod)
def __eq__(self, other):
try:
return super().__eq__(other)
except (IterableUnitCoercionError, UnitOperationError):
return np.zeros(self.shape, dtype="bool")
def __ne__(self, other):
try:
return super().__ne__(other)
except (IterableUnitCoercionError, UnitOperationError):
return np.ones(self.shape, dtype="bool")
#
# Start operation methods
#
def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
func = getattr(ufunc, method)
if "out" not in kwargs:
if ufunc in multiple_output_operators:
out = (None,) * multiple_output_operators[ufunc]
out_func = out
else:
out = None
out_func = None
else:
# we need to get both the actual "out" object and a view onto it
# in case we need to do in-place operations
out = kwargs.pop("out")
if ufunc in multiple_output_operators:
out_func = []
for arr in out:
out_func.append(arr.view(np.ndarray))
out_func = tuple(out_func)
else:
out = out[0]
if out.dtype.kind in ("u", "i"):
new_dtype = "f" + str(out.dtype.itemsize)
float_values = out.astype(new_dtype)
out.dtype = new_dtype
np.copyto(out, float_values)
out_func = out.view(np.ndarray)
if len(inputs) == 1:
# Unary ufuncs
inp = inputs[0]
u = getattr(inp, "units", None)
if u.dimensions is angle and ufunc in trigonometric_operators:
# ensure np.sin(90*degrees) works as expected
inp = inp.in_units("radian").v
# evaluate the ufunc
out_arr = func(np.asarray(inp), out=out_func, **kwargs)
if ufunc in (multiply, divide) and method == "reduce":
mul, unit = _apply_power_mapping(ufunc, u, inp.size, inp.shape, kwargs)
else:
# get unit of result
mul, unit = self._ufunc_registry[ufunc](u)
# use type(self) here so we can support user-defined
# subclasses of unyt_array
ret_class = type(self)
elif len(inputs) == 2:
# binary ufuncs
i0 = inputs[0]
i1 = inputs[1]
if "dask" in sys.modules and isinstance(i1, _dask.array.core.Array):
# need to short circuit all this to handle binary operations
# like unyt_quantity(2,'m') / unyt_dask_array_instance
# only need to check the second argument as if the first arg
# is a unyt_dask_array, it won't end up here.
return i1.__array_ufunc__(ufunc, method, *inputs, **kwargs)
# coerce inputs to be ndarrays if they aren't already
inp0 = _coerce_iterable_units(i0)
inp1 = _coerce_iterable_units(i1)
u0 = getattr(i0, "units", None) or getattr(inp0, "units", None)
u1 = getattr(i1, "units", None) or getattr(inp1, "units", None)
ret_class = _get_binary_op_return_class(type(i0), type(i1))
if u0 is None:
u0 = Unit(registry=getattr(u1, "registry", None))
if u1 is None and ufunc is not power:
u1 = Unit(registry=getattr(u0, "registry", None))
elif ufunc is power:
u1 = inp1
if inp0.shape != () and inp1.shape != ():
raise UnitOperationError(ufunc, u0, u1)
if isinstance(u1, unyt_array):
if u1.units.is_dimensionless:
pass
else:
raise UnitOperationError(ufunc, u0, u1.units)
if u1.shape == ():
u1 = float(u1)
else:
u1 = 1.0
unit_operator = self._ufunc_registry[ufunc]
if (
unit_operator is _preserve_units
and u0.dimensions is temperature
and u1 is not None
and u1.base_offset != 0.0
and u0.base_offset == 0.0
and str(u0.expr) in ["K", "R"]
):
raise UnitOperationError(ufunc, u0, u1)
if unit_operator in (
_preserve_units,
_comparison_unit,
_arctan2_unit,
_difference_units,
):
# check "is" equality first for speed
if u0 is not u1 and u0 != u1:
# we allow adding, multiplying, comparisons with
# zero-filled arrays, lists, etc or scalar zero. We
# do not allow zero-filled unyt_array instances for
# performance reasons. If we did allow it, every
# binary operation would need to scan over all the
# elements of both arrays to check for arrays filled
# with zeros
if not isinstance(i0, unyt_array) or not isinstance(i1, unyt_array):
any_nonzero = [np.count_nonzero(i0), np.count_nonzero(i1)]
if any_nonzero[0] == 0:
u0 = u1
elif any_nonzero[1] == 0:
u1 = u0
if not u0.same_dimensions_as(u1):
if unit_operator is _comparison_unit:
# we allow comparisons between data with
# units and dimensionless data
if u0.is_dimensionless:
u0 = u1
elif u1.is_dimensionless:
u1 = u0
else:
# comparison with different units, so need to check if
# this is == and != which we allow and handle in a
# special way using an early return from __array_ufunc__
if ufunc in (equal, not_equal):
if ufunc is equal:
func = np.zeros_like
else:
func = np.ones_like
ret = func(np.asarray(inp1), dtype=bool)
if out is not None:
out[:] = ret[:]
if isinstance(out, unyt_array):
out.units = Unit(
"", registry=self.units.registry
)
if ret.shape == ():
ret = bool(ret)
return ret
else:
raise UnitOperationError(ufunc, u0, u1)
else:
raise UnitOperationError(ufunc, u0, u1)
conv, offset = u1.get_conversion_factor(u0, inp1.dtype)
new_dtype = np.dtype("f" + str(inp1.dtype.itemsize))
conv = new_dtype.type(conv)
if (
offset is not None
and u1.base_offset != 0.0
and not repr(u0).startswith("delta_")
):
raise InvalidUnitOperation(
"Quantities with units of Fahrenheit or Celsius "
"cannot be multiplied, divided, subtracted or "
"added with data that has different units."
)
inp1 = np.asarray(inp1, dtype=new_dtype) * conv
# get the unit of the result
mul, unit = unit_operator(u0, u1)
# actually evaluate the ufunc
out_arr = func(
inp0.view(np.ndarray), inp1.view(np.ndarray), out=out_func, **kwargs
)
if unit_operator in (_multiply_units, _divide_units):
if unit.is_dimensionless and unit.base_value != 1.0:
if not u0.is_dimensionless:
if u0.dimensions == u1.dimensions:
out_arr = np.multiply(
out_arr.view(np.ndarray), unit.base_value, out=out_func
)
unit = Unit(registry=unit.registry)
if (
u0.base_offset
and u0.dimensions is temperature
or u1.base_offset
and u1.dimensions is temperature
):
raise InvalidUnitOperation(
"Quantities with units of Fahrenheit or Celsius "
"cannot be multiplied, divided, subtracted or added."
)
else:
if ufunc is clip:
inp = []
for i in inputs:
if isinstance(i, unyt_array):
inp.append(i.to(inputs[0].units).view(np.ndarray))
else:
inp.append(i)
if out is not None:
_out = out.view(np.ndarray)
else:
_out = None
out_arr = ufunc(*inp, out=_out)
unit = inputs[0].units
ret_class = type(inputs[0])
mul = 1
else:
raise RuntimeError(
"Support for the %s ufunc with %i inputs has not been "
"added to unyt_array." % (str(ufunc), len(inputs))
)
if unit is None:
out_arr = np.array(out_arr, copy=_COPY_IF_NEEDED)
elif ufunc in (modf, divmod_):
out_arr = tuple(ret_class(o, unit) for o in out_arr)
elif out_arr.shape == ():
out_arr = unyt_quantity(np.asarray(out_arr), unit)
elif out_arr.size == 1:
out_arr = unyt_array(np.asarray(out_arr), unit)
else:
if ret_class is unyt_quantity:
# This happens if you do ndarray * unyt_quantity.
# Explicitly casting to unyt_array avoids creating a
# unyt_quantity with size > 1
out_arr = unyt_array(out_arr, unit)
else:
out_arr = ret_class(out_arr, unit, bypass_validation=True)
if out is not None:
if mul != 1:
multiply(out, mul, out=out)
if np.shares_memory(out_arr, out):
mul = 1
if isinstance(out, unyt_array):
try:
out.units = out_arr.units
except AttributeError:
# out_arr is an ndarray
out.units = Unit("", registry=self.units.registry)
elif isinstance(out, tuple):
for o, oa in zip(out, out_arr):
if o is None:
continue
o.units = oa.units
if mul == 1:
return out_arr
return mul * out_arr
def __array_function__(self, func, types, args, kwargs):
# Follow NEP 18 guidelines
# https://numpy.org/neps/nep-0018-array-function-protocol.html
from unyt._array_functions import _HANDLED_FUNCTIONS, _UNSUPPORTED_FUNCTIONS
if func in _UNSUPPORTED_FUNCTIONS:
# following NEP 18, return NotImplemented as a sentinel value
# which will lead to raising a TypeError, while
# leaving other arguments a chance to take the lead
return NotImplemented
if func not in _HANDLED_FUNCTIONS:
# default to numpy's private implementation
return func._implementation(*args, **kwargs)
# Note: this allows subclasses that don't override
# __array_function__ to handle unyt_array objects
if not all(issubclass(t, unyt_array) or t is np.ndarray for t in types):
return NotImplemented
return _HANDLED_FUNCTIONS[func](*args, **kwargs)
def copy(self, order="C"):
"""
Return a copy of the array.
Parameters
----------
order : {'C', 'F', 'A', 'K'}, optional
Controls the memory layout of the copy. 'C' means C-order,
'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous,
'C' otherwise. 'K' means match the layout of `a` as closely
as possible. (Note that this function and :func:`numpy.copy`
are very similar, but have different default values for their
order= arguments.)
See also
--------
numpy.copy
numpy.copyto
Examples
--------
>>> from unyt import km
>>> x = [[1,2,3],[4,5,6]] * km
>>> y = x.copy()
>>> x.fill(0)
>>> print(x)
[[0 0 0]
[0 0 0]] km
>>> print(y)
[[1 2 3]
[4 5 6]] km
"""
name = getattr(self, "name", None)
try:
return type(self)(np.copy(np.asarray(self)), self.units, name=name)
except TypeError:
# subclasses might not take name as a kwarg
return type(self)(np.copy(np.asarray(self)), self.units)
def __array_finalize__(self, obj):
self.units = getattr(obj, "units", NULL_UNIT)
self.name = getattr(obj, "name", None)
def __pos__(self):
"""Posify the data."""
# this needs to be defined for all numpy versions, see
# numpy issue #9081
return type(self)(super().__pos__(), self.units)
def dot(self, b, out=None):
"""dot product of two arrays.
Refer to `numpy.dot` for full documentation.
See Also
--------
numpy.dot : equivalent function
Examples
--------
>>> from unyt import km, s
>>> a = np.eye(2)*km
>>> b = (np.ones((2, 2)) * 2)*s
>>> print(a.dot(b))
[[2. 2.]
[2. 2.]] km*s
This array method can be conveniently chained:
>>> print(a.dot(b).dot(b))
[[8. 8.]
[8. 8.]] km*s**2
"""
res_units = self.units * getattr(b, "units", NULL_UNIT)
ret = self.view(np.ndarray).dot(np.asarray(b), out=out) * res_units
if out is not None:
out.units = res_units
return ret
def __reduce__(self):
"""Pickle reduction method
See the documentation for the standard library pickle module:
http://docs.python.org/2/library/pickle.html
Unit metadata is encoded in the zeroth element of third element of the
returned tuple, itself a tuple used to restore the state of the
ndarray. This is always defined for numpy arrays.
"""
np_ret = super().__reduce__()
obj_state = np_ret[2]
unit_state = (((str(self.units), self.units.registry.lut),) + obj_state[:],)
new_ret = np_ret[:2] + unit_state + np_ret[3:]
return new_ret
def __setstate__(self, state):
"""Pickle setstate method
This is called inside pickle.read() and restores the unit data from the
metadata extracted in __reduce__ and then serialized by pickle.
"""
super().__setstate__(state[1:])
unit, lut = state[0]
lut = _correct_old_unit_registry(lut)
registry = UnitRegistry(lut=lut, add_default_symbols=False)
self.units = Unit(unit, registry=registry)
def __deepcopy__(self, memodict=None):
"""copy.deepcopy implementation
This is necessary for stdlib deepcopy of arrays and quantities.
"""
ret = super().__deepcopy__(memodict)
try:
return type(self)(ret, copy.deepcopy(self.units), name=self.name)
except TypeError:
# subclasses might not take name as a kwarg
return type(self)(ret, copy.deepcopy(self.units))
class unyt_quantity(unyt_array):
"""
A scalar associated with a unit.
Parameters
----------
input_scalar : an integer or floating point scalar
The scalar to attach units to
units : String unit specification, unit symbol object, or astropy units
The units of the quantity. Powers must be specified using python syntax
(cm**3, not cm^3).
registry : A UnitRegistry object
The registry to create units from. If units is already associated
with a unit registry and this is specified, this will be used instead
of the registry associated with the unit object.
dtype : data-type
The dtype of the array data.
bypass_validation : boolean
If True, all input validation is skipped. Using this option may produce
corrupted, invalid units or array data, but can lead to significant
speedups in the input validation logic adds significant overhead. If
set, units *must* be a valid unit object. Defaults to False.
name : string
The name of the scalar. Defaults to None. This attribute does not propagate
through mathematical operations, but is preserved under indexing
and unit conversions.
Examples
--------
>>> a = unyt_quantity(3., 'cm')
>>> b = unyt_quantity(2., 'm')
>>> print(a + b)
203.0 cm
>>> print(b + a)
2.03 m
NumPy ufuncs will pass through units where appropriate.
>>> import numpy as np
>>> from unyt import g, cm
>>> a = 12*g/cm**3
>>> print(np.abs(a))
12 g/cm**3
and strip them when it would be annoying to deal with them.
>>> print(np.log10(a))
1.0791812460476249
"""
def __new__(
cls,
input_scalar,
units=None,
registry=None,
dtype=None,
*,
bypass_validation=False,
name=None,
):
input_units = units
if not (
bypass_validation
or isinstance(input_scalar, (numeric_type, np.number, np.ndarray))
):
raise RuntimeError("unyt_quantity values must be numeric")
if input_units is None:
units = getattr(input_scalar, "units", None)
else:
units = input_units
ret = unyt_array.__new__(
cls,
np.asarray(input_scalar),
units,
registry,
dtype=dtype,
bypass_validation=bypass_validation,
name=name,
)
if ret.size > 1:
raise RuntimeError("unyt_quantity instances must be scalars")
return ret
def __round__(self):
return type(self)(round(float(self)), self.units)
def reshape(self, *shape, order="C"):
# this is necessary to support some numpy operations
# natively, like numpy.meshgrid, which internally performs
# reshaping, e.g., arr.reshape(1, -1), which doesn't affect the size,
# but does change the object's internal representation to a >0D array
# see https://github.com/yt-project/unyt/issues/224
if len(shape) == 1:
shape = shape[0]
if shape == () or shape is None:
return super().reshape(shape, order=order)
else:
return unyt_array(self).reshape(shape, order=order)
def _validate_numpy_wrapper_units(v, arrs):
if not any(isinstance(a, unyt_array) for a in arrs):
return v
if not all(isinstance(a, unyt_array) for a in arrs):
raise RuntimeError("Not all of your arrays are unyt_arrays.")
a1 = arrs[0]
if not all(a.units == a1.units for a in arrs[1:]):
raise RuntimeError("Your arrays must have identical units.")
v.units = a1.units
return v
def uconcatenate(arrs, axis=0):
"""Concatenate a sequence of arrays.
This wrapper around numpy.concatenate preserves units. All input arrays
must have the same units. See the documentation of numpy.concatenate for
full details.
Examples
--------
>>> from unyt import cm
>>> A = [1, 2, 3]*cm
>>> B = [2, 3, 4]*cm
>>> uconcatenate((A, B)) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
unyt_array([1, 2, 3, 2, 3, 4], 'cm')
"""
warn_deprecated(
"unyt.uconcatenate", replacement="use numpy.concatenate", since_version="3.0"
)
v = np.concatenate._implementation(arrs, axis=axis)
v = _validate_numpy_wrapper_units(v, arrs)
return v
def ucross(arr1, arr2, registry=None, axisa=-1, axisb=-1, axisc=-1, axis=None):
"""Applies the cross product to two YT arrays.
This wrapper around numpy.cross preserves units.
See the documentation of numpy.cross for full
details.
"""
warn_deprecated(
"unyt.ucross",
replacement=(
"use numpy.cross (note that the *registry* argument will not be ported)"
),
since_version="3.0",
)
v = np.cross._implementation(
arr1, arr2, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis
)
units = arr1.units * arr2.units
arr = unyt_array(v, units, registry=registry)
return arr
def uintersect1d(arr1, arr2, assume_unique=False):
"""Find the sorted unique elements of the two input arrays.
A wrapper around numpy.intersect1d that preserves units. All input arrays
must have the same units. See the documentation of numpy.intersect1d for
full details.
Examples
--------
>>> from unyt import cm
>>> A = [1, 2, 3]*cm
>>> B = [2, 3, 4]*cm
>>> uintersect1d(A, B) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
unyt_array([2, 3], 'cm')
"""
warn_deprecated(
"unyt.uintersect1d", replacement="use numpy.intersect1d", since_version="3.0"
)
return np.intersect1d(arr1, arr2, assume_unique=assume_unique)
def uunion1d(arr1, arr2):
"""Find the union of two arrays.
A wrapper around numpy.intersect1d that preserves units. All input arrays
must have the same units. See the documentation of numpy.intersect1d for
full details.
Examples
--------
>>> from unyt import cm
>>> A = [1, 2, 3]*cm
>>> B = [2, 3, 4]*cm
>>> uunion1d(A, B) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
unyt_array([1, 2, 3, 4], 'cm')
"""
warn_deprecated(
"unyt.uunion1d", replacement="use numpy.union1d", since_version="3.0"
)
return np.union1d(arr1, arr2)
def unorm(data, ord=None, axis=None, keepdims=False):
"""Matrix or vector norm that preserves units
This is a wrapper around np.linalg.norm that preserves units. See
the documentation for that function for descriptions of the keyword
arguments.
Examples
--------
>>> from unyt import km
>>> data = [1, 2, 3]*km
>>> print(unorm(data)) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
3.7416573867739413 km
"""
warn_deprecated(
"unyt.unorm", replacement="use numpy.linalg.norm", since_version="3.0"
)
norm = np.linalg.norm._implementation(data, ord=ord, axis=axis, keepdims=keepdims)
if norm.shape == ():
return unyt_quantity(norm, data.units)
return unyt_array(norm, data.units)
def udot(op1, op2):
"""Matrix or vector dot product that preserves units
This is a wrapper around np.dot that preserves units.
Examples
--------
>>> from unyt import km, s
>>> a = np.eye(2)*km
>>> b = (np.ones((2, 2)) * 2)*s
>>> print(udot(a, b)) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
[[2. 2.]
[2. 2.]] km*s
"""
warn_deprecated("unyt.udot", replacement="use numpy.dot", since_version="3.0")
dot = np.dot._implementation(op1.d, op2.d)
units = op1.units * op2.units
if dot.shape == ():
return unyt_quantity(dot, units)
return unyt_array(dot, units)
def uvstack(arrs):
"""Stack arrays in sequence vertically (row wise) while preserving units
This is a wrapper around np.vstack that preserves units.
Examples
--------
>>> from unyt import km
>>> a = [1, 2, 3]*km
>>> b = [2, 3, 4]*km
>>> print(uvstack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
[[1 2 3]
[2 3 4]] km
"""
warn_deprecated("unyt.uvstack", replacement="use numpy.vstack", since_version="3.0")
v = np.vstack._implementation(arrs)
v = _validate_numpy_wrapper_units(v, arrs)
return v
def uhstack(arrs):
"""Stack arrays in sequence horizontally while preserving units
This is a wrapper around np.hstack that preserves units.
Examples
--------
>>> from unyt import km
>>> a = [1, 2, 3]*km
>>> b = [2, 3, 4]*km
>>> print(uhstack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
[1 2 3 2 3 4] km
>>> a = [[1],[2],[3]]*km
>>> b = [[2],[3],[4]]*km
>>> print(uhstack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
[[1 2]
[2 3]
[3 4]] km
"""
warn_deprecated("unyt.uhstack", replacement="use numpy.hstack", since_version="3.0")
v = np.hstack._implementation(arrs)
v = _validate_numpy_wrapper_units(v, arrs)
return v
def ustack(arrs, axis=0):
"""Join a sequence of arrays along a new axis while preserving units
The axis parameter specifies the index of the new axis in the
dimensions of the result. For example, if ``axis=0`` it will be the
first dimension and if ``axis=-1`` it will be the last dimension.
This is a wrapper around np.stack that preserves units. See the
documentation for np.stack for full details.
Examples
--------
>>> from unyt import km
>>> a = [1, 2, 3]*km
>>> b = [2, 3, 4]*km
>>> print(ustack([a, b])) # doctest: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
DeprecationWarning: ...
[[1 2 3]
[2 3 4]] km
"""
warn_deprecated("unyt.ustack", replacement="use numpy.stack", since_version="3.0")
v = np.stack._implementation(arrs, axis=axis)
v = _validate_numpy_wrapper_units(v, arrs)
return v
def _get_binary_op_return_class(cls1, cls2):
if cls1 is cls2:
return cls1
if cls1 in (Unit, np.ndarray, np.matrix, np.ma.masked_array) or issubclass(
cls1, (numeric_type, np.number, list, tuple)
):
return cls2
if cls2 in (Unit, np.ndarray, np.matrix, np.ma.masked_array) or issubclass(
cls2, (numeric_type, np.number, list, tuple)
):
return cls1
if issubclass(cls1, unyt_quantity):
return cls2
if issubclass(cls2, unyt_quantity):
return cls1
if issubclass(cls1, cls2):
return cls1
if issubclass(cls2, cls1):
return cls2
else:
raise RuntimeError(
"Undefined operation for a unyt_array subclass. "
f"Received operand types ({cls1}) and ({cls2})"
)
def loadtxt(fname, dtype="float", delimiter="\t", usecols=None, comments="#"):
r"""
Load unyt_arrays with unit information from a text file. Each row in the
text file must have the same number of values.
Parameters
----------
fname : str
Filename to read.
dtype : data-type, optional
Data-type of the resulting array; default: float.
delimiter : str, optional
The string used to separate values. By default, this is any
whitespace.
usecols : sequence, optional
Which columns to read, with 0 being the first. For example,
``usecols = (1,4,5)`` will extract the 2nd, 5th and 6th columns.
The default, None, results in all columns being read.
comments : str, optional
The character used to indicate the start of a comment;
default: '#'.
Examples
--------
>>> temp, velx = loadtxt(
... "sphere.dat", usecols=(1,2), delimiter="\t") # doctest: +SKIP
"""
f = open(fname)
next_one = False
units = []
num_cols = -1
for line in f.readlines():
words = line.strip().split()
if len(words) == 0:
continue
if line[0] == comments:
if next_one:
units = words[1:]
if len(words) == 2 and words[1] == "Units":
next_one = True
else:
# Here we catch the first line of numbers
col_words = line.strip().split(delimiter)
for word in col_words:
# test that word can be converted to a number
complex(word)
num_cols = len(col_words)
break
f.close()
if len(units) != num_cols:
units = ["dimensionless"] * num_cols
arrays = np.loadtxt(
fname,
dtype=dtype,
comments=comments,
delimiter=delimiter,
converters=None,
unpack=True,
usecols=usecols,
ndmin=0,
)
if len(arrays.shape) < 2:
arrays = [arrays]
if usecols is not None:
units = [units[col] for col in usecols]
ret = tuple(unyt_array(arr, unit) for arr, unit in zip(arrays, units))
if len(ret) == 1:
return ret[0]
return ret
def savetxt(
fname, arrays, fmt="%.18e", delimiter="\t", header="", footer="", comments="#"
):
r"""
Write unyt_arrays with unit information to a text file.
Parameters
----------
fname : str
The file to write the unyt_arrays to.
arrays : list of unyt_arrays or single unyt_array
The array(s) to write to the file.
fmt : str or sequence of strs, optional
A single format (%10.5f), or a sequence of formats.
delimiter : str, optional
String or character separating columns.
header : str, optional
String that will be written at the beginning of the file, before the
unit header.
footer : str, optional
String that will be written at the end of the file.
comments : str, optional
String that will be prepended to the ``header`` and ``footer`` strings,
to mark them as comments. Default: '# ', as expected by e.g.
``unyt.loadtxt``.
Examples
--------
>>> import unyt as u
>>> a = [1, 2, 3]*u.cm
>>> b = [8, 10, 12]*u.cm/u.s
>>> c = [2, 85, 9]*u.g
>>> savetxt("sphere.dat", [a,b,c], header='My sphere stuff',
... delimiter="\t") # doctest: +SKIP
"""
if not isinstance(arrays, list):
arrays = [arrays]
units = []
for array in arrays:
if hasattr(array, "units"):
units.append(str(array.units))
else:
units.append("dimensionless")
if header != "" and not header.endswith("\n"):
header += "\n"
header += " Units\n " + "\t".join(units)
np.savetxt(
fname,
np.transpose(arrays),
header=header,
fmt=fmt,
delimiter=delimiter,
footer=footer,
newline="\n",
comments=comments,
)
def allclose_units(actual, desired, rtol=1e-7, atol=0, **kwargs):
"""Returns False if two objects are not equal up to desired tolerance
This is a wrapper for :func:`numpy.allclose` that also
verifies unit consistency
Parameters
----------
actual : array-like
Array obtained (possibly with attached units)
desired : array-like
Array to compare with (possibly with attached units)
rtol : float, optional
Relative tolerance, defaults to 1e-7
atol : float or quantity, optional
Absolute tolerance. If units are attached, they must be consistent
with the units of ``actual`` and ``desired``. If no units are attached,
assumes the same units as ``desired``. Defaults to zero.
Raises
------
RuntimeError
If units of ``rtol`` are not dimensionless
See Also
--------
:func:`unyt.testing.assert_allclose_units`
Notes
-----
Also accepts additional keyword arguments accepted by
:func:`numpy.allclose`, see the documentation of that
function for details.
Examples
--------
>>> import unyt as u
>>> actual = [1e-5, 1e-3, 1e-1]*u.m
>>> desired = actual.to("cm")
>>> allclose_units(actual, desired)
True
"""
# Create a copy to ensure this function does not alter input arrays
act = unyt_array(actual)
des = unyt_array(desired)
try:
des = des.in_units(act.units)
except (UnitOperationError, UnitConversionError):
return False
rt = unyt_array(rtol)
if not rt.units.is_dimensionless:
raise RuntimeError(f"Units of rtol ({rt.units}) are not dimensionless")
if not isinstance(atol, unyt_array):
at = unyt_quantity(atol, des.units)
else:
at = atol
try:
at = at.in_units(act.units)
except (UnitOperationError, UnitConversionError):
return False
# units have been validated, so we strip units before calling numpy
# to avoid spurious errors
act = act.value
des = des.value
rt = rt.value
at = at.value
return np.allclose(act, des, rt, at, **kwargs)