1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
import pytest
from pint import UnitRegistry
# Conditionally import NumPy and any upcast type libraries
np = pytest.importorskip("numpy", reason="NumPy is not available")
sparse = pytest.importorskip("sparse", reason="sparse is not available")
da = pytest.importorskip("dask.array", reason="Dask is not available")
# Set up unit registry and sample
ureg = UnitRegistry(force_ndarray_like=True)
q_base = (np.arange(25).reshape(5, 5).T + 1) * ureg.kg
# Define identity function for use in tests
def identity(x):
return x
@pytest.fixture(params=["sparse", "masked_array", "dask_array"])
def array(request):
"""Generate 5x5 arrays of given type for tests."""
if request.param == "sparse":
# Create sample sparse COO as a permutation matrix.
coords = [[0, 1, 2, 3, 4], [1, 3, 0, 2, 4]]
data = [1.0] * 5
return sparse.COO(coords, data, shape=(5, 5))
elif request.param == "masked_array":
# Create sample masked array as an upper triangular matrix.
return np.ma.masked_array(
np.arange(25, dtype=np.float).reshape((5, 5)),
mask=np.logical_not(np.triu(np.ones((5, 5)))),
)
elif request.param == "dask_array":
return da.arange(25, chunks=5, dtype=float).reshape((5, 5))
@pytest.mark.parametrize(
"op, magnitude_op, unit_op",
[
pytest.param(identity, identity, identity, id="identity"),
pytest.param(
lambda x: x + 1 * ureg.m, lambda x: x + 1, identity, id="addition"
),
pytest.param(
lambda x: x - 20 * ureg.cm, lambda x: x - 0.2, identity, id="subtraction"
),
pytest.param(
lambda x: x * (2 * ureg.s),
lambda x: 2 * x,
lambda u: u * ureg.s,
id="multiplication",
),
pytest.param(
lambda x: x / (1 * ureg.s), identity, lambda u: u / ureg.s, id="division"
),
pytest.param(lambda x: x ** 2, lambda x: x ** 2, lambda u: u ** 2, id="square"),
pytest.param(lambda x: x.T, lambda x: x.T, identity, id="transpose"),
pytest.param(np.mean, np.mean, identity, id="mean ufunc"),
pytest.param(np.sum, np.sum, identity, id="sum ufunc"),
pytest.param(np.sqrt, np.sqrt, lambda u: u ** 0.5, id="sqrt ufunc"),
pytest.param(
lambda x: np.reshape(x, (25,)),
lambda x: np.reshape(x, (25,)),
identity,
id="reshape function",
),
pytest.param(np.amax, np.amax, identity, id="amax function"),
],
)
def test_univariate_op_consistency(op, magnitude_op, unit_op, array):
q = ureg.Quantity(array, "meter")
res = op(q)
assert np.all(res.magnitude == magnitude_op(array)) # Magnitude check
assert res.units == unit_op(q.units) # Unit check
assert q.magnitude is array # Immutability check
@pytest.mark.parametrize(
"op, unit",
[
pytest.param(lambda x, y: x * y, ureg("kg m"), id="multiplication"),
pytest.param(lambda x, y: x / y, ureg("m / kg"), id="division"),
pytest.param(np.multiply, ureg("kg m"), id="multiply ufunc"),
],
)
def test_bivariate_op_consistency(op, unit, array):
q = ureg.Quantity(array, "meter")
res = op(q, q_base)
assert np.all(res.magnitude == op(array, q_base.magnitude)) # Magnitude check
assert res.units == unit # Unit check
assert q.magnitude is array # Immutability check
@pytest.mark.parametrize(
"op",
[
pytest.param(
lambda a, u: a * u,
id="array-first",
marks=pytest.mark.xfail(reason="upstream issue numpy/numpy#15200"),
),
pytest.param(lambda a, u: u * a, id="unit-first"),
],
)
@pytest.mark.parametrize(
"unit",
[pytest.param(ureg.m, id="Unit"), pytest.param(ureg("meter"), id="Quantity")],
)
def test_array_quantity_creation_by_multiplication(op, unit, array):
assert type(op(array, unit)) == ureg.Quantity
|
