Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP] Addition of UnitOperationWarning and strict kwarg to UnitRegistry #166

Open
wants to merge 20 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 17 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions docs/usage.rst
Original file line number Diff line number Diff line change
Expand Up @@ -1024,6 +1024,20 @@ When writing tests, it is convenient to use :mod:`unyt.testing`. In particular,
>>> desired = actual.to("cm")
>>> assert_allclose_units(actual, desired)

Integrating :mod:`unyt` Into a Legacy Code with Non-Strict Mode
---------------------------------------------------------------

If using a custom :class:`UnitRegistry <unyt.unit_registry.UnitRegistry>`, it
is possible to supply ``strict=False`` when initializing. This will change the behavior
when an invalid operation is attempted. Instead of raising a :class:`UnitOperationError <unyt.exceptions.UnitOperationError>`,
a :class:`UnitOperationWarning <unyt.exceptions.UnitOperationError>` will be
provided instead, units will be stripped, and the operation will return a value
without units.

This behavior is not recommended in general since many of the important benefits
of :mod:`unyt` are lost, but it can be useful for integrating :mod:`unyt` into
legacy code which was not previously unit-aware so that unit support can be
added gradually.

Custom Unit Systems
-------------------
Expand Down
35 changes: 29 additions & 6 deletions unyt/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,10 @@
# -----------------------------------------------------------------------------

import copy

import re
from functools import lru_cache
from numbers import Number as numeric_type
import re

import numpy as np
from numpy import (
add,
Expand Down Expand Up @@ -123,6 +123,7 @@
MKSCGSConversionError,
UnitOperationError,
UnitConversionError,
UnitOperationWarning,
UnitsNotReducible,
SymbolNotFoundError,
)
Expand Down Expand Up @@ -166,6 +167,20 @@ def _iterable(obj):
return True


def _unit_operation_error_raise_or_warn(ufunc, u0, u1, func, *inputs):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am hesitant about how this function implements the proposed functionality. Specifically, it feels weird to replace a raise statement with a return statement (where this function is used).
Another comment is that the func argument is redundant (I think), since it can be obtained as func = ufunc.method.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I share this uneasiness but I'm not sure of a better solution. It seems clear to me that some users will want this to fail explicitly and raise, whereas others will just want the warning, so there has to be a conditional somewhere that switches between that behaviour?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To your second comment, makes sense, I'll remove the func argument.

if (
hasattr(u0, "units")
and not u0.units.registry.strict
and hasattr(u1, "units")
and not u1.units.registry.strict
):
warnings.warn(UnitOperationWarning(ufunc, u0, u1))
unwrapped_inputs = [i.value if isinstance(i, unyt_array) else i for i in inputs]
return func(*unwrapped_inputs)
else:
raise UnitOperationError(ufunc, u0, u1)


@lru_cache(maxsize=128, typed=False)
def _sqrt_unit(unit):
return 1, unit ** 0.5
Expand Down Expand Up @@ -1759,12 +1774,16 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
elif ufunc is power:
u1 = inp1
if inp0.shape != () and inp1.shape != ():
raise UnitOperationError(ufunc, u0, u1)
return _unit_operation_error_raise_or_warn(
ufunc, u0, u1, func, *inputs
)
if isinstance(u1, unyt_array):
if u1.units.is_dimensionless:
pass
else:
raise UnitOperationError(ufunc, u0, u1.units)
return _unit_operation_error_raise_or_warn(
ufunc, u0, u1.units, func, *inputs
)
if u1.shape == ():
u1 = float(u1)
else:
Expand Down Expand Up @@ -1814,9 +1833,13 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
ret = bool(ret)
return ret
else:
raise UnitOperationError(ufunc, u0, u1)
return _unit_operation_error_raise_or_warn(
ufunc, u0, u1, func, *inputs
)
else:
raise UnitOperationError(ufunc, u0, u1)
return _unit_operation_error_raise_or_warn(
ufunc, u0, u1, func, *inputs
)
conv, offset = u1.get_conversion_factor(u0, inp1.dtype)
new_dtype = np.dtype("f" + str(inp1.dtype.itemsize))
conv = new_dtype.type(conv)
Expand Down
23 changes: 23 additions & 0 deletions unyt/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,29 @@ def __str__(self):
return err


class UnitOperationWarning(UserWarning):
"""A warning that is raised when unit operations are not allowed but
when running with a UnitRegistry with strict=False. In this case,
the operation is allowed to continue but with unit information stripped.
"""

def __init__(self, operation, unit1, unit2=None):
self.operation = operation
self.unit1 = unit1
self.unit2 = unit2
UserWarning.__init__(self)

def __str__(self):
err = (
f"The {self.operation.__name__} operator for unyt_arrays"
f" with units {self.unit1!r} (dimensions {self.unit1.dimensions!r})"
)
if self.unit2 is not None:
err += f" and {self.units2!r} (dimensions {self.units2.dimenstions!r})"
err += " is not well defined. Performing operation without units instead."
return err


class UnitConversionError(Exception):
"""An error raised when converting to a unit with different dimensions.

Expand Down
27 changes: 20 additions & 7 deletions unyt/tests/test_unyt_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,22 +17,25 @@
import copy
import itertools
import math
import numpy as np
import operator
import os
import pickle
import pytest
import shutil
import tempfile
import warnings

import numpy as np
import pytest
from numpy import array
from numpy.testing import (
assert_array_equal,
assert_equal,
assert_array_almost_equal,
assert_almost_equal,
)
from numpy import array
from unyt import dimensions, Unit, degC, K, delta_degC, degF, R, delta_degF
from unyt._on_demand_imports import _astropy, _h5py, _pint, NotAModule
from unyt._physical_ratios import metallicity_sun, speed_of_light_cm_per_s
from unyt.array import (
unyt_array,
unyt_quantity,
Expand All @@ -55,15 +58,13 @@
IterableUnitCoercionError,
UnitConversionError,
UnitOperationError,
UnitOperationWarning,
UnitParseError,
UnitsNotReducible,
)
from unyt.testing import assert_allclose_units, _process_warning
from unyt.unit_symbols import cm, m, g, degree
from unyt.unit_registry import UnitRegistry
from unyt._on_demand_imports import _astropy, _h5py, _pint, NotAModule
from unyt._physical_ratios import metallicity_sun, speed_of_light_cm_per_s
from unyt import dimensions, Unit, degC, K, delta_degC, degF, R, delta_degF
from unyt.unit_symbols import cm, m, g, degree


def operate_and_compare(a, b, op, answer):
Expand Down Expand Up @@ -2539,3 +2540,15 @@ def test_invalid_unit_quantity_from_string(s):
match="Could not find unit symbol '{}' in the provided symbols.".format(un_str),
):
unyt_quantity.from_string(s)


def test_non_strict_registry():

reg = UnitRegistry(strict=False)

a1 = unyt_array([1, 2, 3], "m", registry=reg)
a2 = unyt_array([4, 5, 6], "kg", registry=reg)

with pytest.warns(UnitOperationWarning):
answer = operator.add(a1, a2)
assert_array_equal(answer, [5, 7, 9])
10 changes: 9 additions & 1 deletion unyt/unit_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,9 @@ class UnitRegistry:

_unit_system_id = None

def __init__(self, add_default_symbols=True, lut=None, unit_system=None):
def __init__(
self, add_default_symbols=True, lut=None, unit_system=None, strict=True
mkhorton marked this conversation as resolved.
Show resolved Hide resolved
):
self._unit_object_cache = {}
if lut:
self.lut = lut
Expand All @@ -67,6 +69,12 @@ def __init__(self, add_default_symbols=True, lut=None, unit_system=None):
if add_default_symbols:
self.lut.update(default_unit_symbol_lut)

# This boolean determines whether to raise a UnitOperationError or
# strip units and provide a UnitOperationWarning if an invalid
# operation is attempted. The default is strict=True and is
# strongly recommended.
self.strict = strict

def __getitem__(self, key):
try:
ret = self.lut[str(key)]
Expand Down