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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
|
from __future__ import annotations
import operator
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")
def WR(func):
"""Function to wrap another containing 1 argument.
Used to parametrize tests in which some cases depend
on the registry while avoiding to create it at the module level
"""
return lambda ureg, x: func(x)
def WR2(func):
"""Function to wrap another containing 2 argument.
Used to parametrize tests in which some cases depend
on the registry while avoiding to create it at the module level
"""
return lambda ureg, x, y: func(x, y)
@pytest.fixture(scope="module")
def local_registry():
# Set up unit registry and sample
return UnitRegistry(force_ndarray_like=True)
@pytest.fixture(scope="module")
def q_base(local_registry):
# Set up unit registry and sample
return (np.arange(25).reshape(5, 5).T + 1) * local_registry.kg
# Define identity function for use in tests
def id_matrix(ureg, 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=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(id_matrix, id_matrix, id_matrix, id="identity"),
pytest.param(
lambda ureg, x: x + 1 * ureg.m,
lambda ureg, x: x + 1,
id_matrix,
id="addition",
),
pytest.param(
lambda ureg, x: x - 20 * ureg.cm,
lambda ureg, x: x - 0.2,
id_matrix,
id="subtraction",
),
pytest.param(
lambda ureg, x: x * (2 * ureg.s),
lambda ureg, x: 2 * x,
lambda ureg, u: u * ureg.s,
id="multiplication",
),
pytest.param(
lambda ureg, x: x / (1 * ureg.s),
id_matrix,
lambda ureg, u: u / ureg.s,
id="division",
),
pytest.param(
WR(lambda x: x**2),
WR(lambda x: x**2),
WR(lambda u: u**2),
id="square",
),
pytest.param(WR(lambda x: x.T), WR(lambda x: x.T), id_matrix, id="transpose"),
pytest.param(WR(np.mean), WR(np.mean), id_matrix, id="mean ufunc"),
pytest.param(WR(np.sum), WR(np.sum), id_matrix, id="sum ufunc"),
pytest.param(WR(np.sqrt), WR(np.sqrt), WR(lambda u: u**0.5), id="sqrt ufunc"),
pytest.param(
WR(lambda x: np.reshape(x, (25,))),
WR(lambda x: np.reshape(x, (25,))),
id_matrix,
id="reshape function",
),
pytest.param(WR(np.amax), WR(np.amax), id_matrix, id="amax function"),
],
)
def test_univariate_op_consistency(
local_registry, q_base, op, magnitude_op, unit_op, array
):
q = local_registry.Quantity(array, "meter")
res = op(local_registry, q)
assert np.all(
res.magnitude == magnitude_op(local_registry, array)
) # Magnitude check
assert res.units == unit_op(local_registry, q.units) # Unit check
assert q.magnitude is array # Immutability check
@pytest.mark.parametrize(
"op, unit",
[
pytest.param(operator.mul, lambda ureg: ureg("kg m"), id="multiplication"),
pytest.param(operator.truediv, lambda ureg: ureg("m / kg"), id="division"),
pytest.param(np.multiply, lambda ureg: ureg("kg m"), id="multiply ufunc"),
],
)
def test_bivariate_op_consistency(local_registry, q_base, op, unit, array):
# This is to avoid having a ureg built at the module level.
unit = unit(local_registry)
q = local_registry.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(
WR2(operator.mul),
id="array-first",
marks=pytest.mark.xfail(reason="upstream issue numpy/numpy#15200"),
),
pytest.param(
WR2(operator.mul),
id="unit-first",
marks=pytest.mark.xfail(reason="upstream issue numpy/numpy#15200"),
),
],
)
@pytest.mark.parametrize(
"unit",
[
pytest.param(lambda ureg: ureg.m, id="Unit"),
pytest.param(lambda ureg: ureg("meter"), id="Quantity"),
],
)
def test_array_quantity_creation_by_multiplication(
local_registry, q_base, op, unit, array
):
# This is to avoid having a ureg built at the module level.
unit = unit(local_registry)
assert type(op(local_registry, array, unit)) == local_registry.Quantity
|
