Source code for unyt.array

"""
unyt_array class.



"""


import copy
import re
import sys
import warnings
from functools import lru_cache
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,
    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 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"],
}

# 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 _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,
    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}


[docs] 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, 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 #
[docs] 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)
[docs] 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, )
[docs] 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 )
[docs] 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)
[docs] 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)
[docs] 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)
[docs] 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
[docs] 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)
[docs] 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")
[docs] 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")
[docs] 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)
[docs] 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)
[docs] 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()
[docs] 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)
[docs] 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)
[docs] 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)
[docs] 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] >>> print(data.argsort()) [0 2 1] """ return self.view(np.ndarray).argsort(axis, kind, order)
[docs] @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)
[docs] 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)
[docs] @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)
[docs] 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)
[docs] @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)
[docs] def to_string(self): # this is implemented purely for symmetry's sake return str(self)
# # End unit conversion methods #
[docs] 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()
[docs] @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=False) 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)
[docs] 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)
[docs] 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))
[docs] 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)
[docs] 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
[docs] 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
[docs] 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
[docs] 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)
[docs] 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)
[docs] 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)
[docs] 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)
[docs] 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
[docs] 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
[docs] 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})" )
[docs] 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
[docs] 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, )
[docs] 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)