From ef91a23d2f3a8a7b77265a005b1947da00d1bd00 Mon Sep 17 00:00:00 2001 From: Andrew Date: Thu, 26 Oct 2023 19:33:02 +0100 Subject: [PATCH] quantity checking functions --- pint/numpy_util.py | 97 +++++++++++++++++++++++++++++++ pint/testsuite/test_numpy_util.py | 81 ++++++++++++++++++++++++++ 2 files changed, 178 insertions(+) create mode 100644 pint/numpy_util.py create mode 100644 pint/testsuite/test_numpy_util.py diff --git a/pint/numpy_util.py b/pint/numpy_util.py new file mode 100644 index 000000000..e70845a48 --- /dev/null +++ b/pint/numpy_util.py @@ -0,0 +1,97 @@ +from .util import iterable, sized +from .compat import np + + +def is_quantity_with_scalar_magnitude(obj): + """Test for Quantity with scalar magnitude. + + Parameters + ---------- + obj : object + + + Returns + ------- + True if obj is a Quantity with a scalar magnitude; False otherwise + """ + return is_quantity(obj) and not iterable(obj._magnitude) + + +def is_quantity_with_sequence_magnitude(obj): + """Test for Quantity with sequence magnitude. + + Parameters + ---------- + obj : object + + + Returns + ------- + True if obj is a Quantity with a sequence magnitude; False otherwise + + Examples + -------- + + >>> is_quantity_with_sequence_magnitude([1, 2, 3]) + False + + >>> is_quantity_with_sequence_magnitude([1, Q_(2, 'm'), 3]) + False + + >>> is_quantity_with_sequence_magnitude(Q_([1, 2, 3], 'm')) + True + """ + return is_quantity(obj) and iterable(obj._magnitude) + + +def is_sequence_with_quantity_elements(obj): + """Test for sequences of quantities. + + Parameters + ---------- + obj : object + + + Returns + ------- + True if obj is a sequence and at least one element is a Quantity; False otherwise + + Examples + -------- + + >>> is_sequence_with_quantity_elements([1, 2, 3]) + False + + >>> is_sequence_with_quantity_elements([1, Q_(2, 'm'), 3]) + True + + >>> is_sequence_with_quantity_elements(Q_([1, 2, 3], 'm')) + True + """ + if np is not None and isinstance(obj, np.ndarray) and not obj.dtype.hasobject: + # If obj is a numpy array, avoid looping on all elements + # if dtype does not have objects + return False + return ( + iterable(obj) + and sized(obj) + and not isinstance(obj, str) + and any(is_quantity(item) for item in obj) + ) + + +def is_quantity(obj): + """Test for _units and _magnitude attrs. + + This is done in place of isinstance(Quantity, arg), which would cause a circular import. + + Parameters + ---------- + obj : Object + + + Returns + ------- + bool + """ + return hasattr(obj, "_units") and hasattr(obj, "_magnitude") diff --git a/pint/testsuite/test_numpy_util.py b/pint/testsuite/test_numpy_util.py new file mode 100644 index 000000000..5346b31f9 --- /dev/null +++ b/pint/testsuite/test_numpy_util.py @@ -0,0 +1,81 @@ +from pint.numpy_util import ( + is_quantity, + is_quantity_with_scalar_magnitude, + is_quantity_with_sequence_magnitude, + is_sequence_with_quantity_elements, +) +from pint.compat import np +import pytest +from pint import Quantity as Q_ + + +@pytest.mark.parametrize( + "obj,result", + [ + (Q_(1, "m"), True), + (Q_(np.nan, "m"), True), + (Q_([1, 2], "m"), True), + (Q_([1, np.nan], "m"), True), + (Q_(np.array([1, 2]), "m"), True), + (Q_(np.array([1, np.nan]), "m"), True), + (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), False), + (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), False), + (np.array([Q_(1, "m"), np.nan], dtype="object"), False), + ], +) +def test_is_quantity(obj, result): + assert is_quantity(obj) == result + + +@pytest.mark.parametrize( + "obj,result", + [ + (Q_(1, "m"), True), + (Q_(np.nan, "m"), True), + (Q_([1, 2], "m"), False), + (Q_([1, np.nan], "m"), False), + (Q_(np.array([1, 2]), "m"), False), + (Q_(np.array([1, np.nan]), "m"), False), + (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), False), + (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), False), + (np.array([Q_(1, "m"), np.nan], dtype="object"), False), + ], +) +def test_is_quantity_with_scalar_magnitude(obj, result): + assert is_quantity_with_scalar_magnitude(obj) == result + + +@pytest.mark.parametrize( + "obj,result", + [ + (Q_(1, "m"), False), + (Q_(np.nan, "m"), False), + (Q_([1, 2], "m"), True), + (Q_([1, np.nan], "m"), True), + (Q_(np.array([1, 2]), "m"), True), + (Q_(np.array([1, np.nan]), "m"), True), + (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), False), + (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), False), + (np.array([Q_(1, "m"), np.nan], dtype="object"), False), + ], +) +def test_is_quantity_with_sequence_magnitude(obj, result): + assert is_quantity_with_sequence_magnitude(obj) == result + + +@pytest.mark.parametrize( + "obj,result", + [ + (Q_(1, "m"), False), + (Q_(np.nan, "m"), False), + (Q_([1, 2], "m"), True), + (Q_([1, np.nan], "m"), True), + (Q_(np.array([1, 2]), "m"), True), + (Q_(np.array([1, np.nan]), "m"), True), + (np.array([Q_(1, "m"), Q_(2, "m")], dtype="object"), True), + (np.array([Q_(1, "m"), Q_(np.nan, "m")], dtype="object"), True), + (np.array([Q_(1, "m"), np.nan], dtype="object"), True), + ], +) +def test_is_sequence_with_quantity_elements(obj, result): + assert is_sequence_with_quantity_elements(obj) == result