codekingpro/portable-devtools

1"""Miscellaneous functions for testing masked arrays and subclasses
2
3:author: Pierre Gerard-Marchant
4:contact: pierregm_at_uga_dot_edu
5
6"""
7import operator
8
9import numpy as np
10import numpy._core.umath as umath
11import numpy.testing
12from numpy import ndarray
13from numpy.testing import (  # noqa: F401
14    assert_,
15    assert_allclose,
16    assert_array_almost_equal_nulp,
17    assert_raises,
18    build_err_msg,
19)
20
21from .core import filled, getmask, mask_or, masked, masked_array, nomask
22
23__all__masked = [
24    'almost', 'approx', 'assert_almost_equal', 'assert_array_almost_equal',
25    'assert_array_approx_equal', 'assert_array_compare',
26    'assert_array_equal', 'assert_array_less', 'assert_close',
27    'assert_equal', 'assert_equal_records', 'assert_mask_equal',
28    'assert_not_equal', 'fail_if_array_equal',
29    ]
30
31# Include some normal test functions to avoid breaking other projects who
32# have mistakenly included them from this file. SciPy is one. That is
33# unfortunate, as some of these functions are not intended to work with
34# masked arrays. But there was no way to tell before.
35from unittest import TestCase  # noqa: F401
36
37__some__from_testing = [
38    'TestCase', 'assert_', 'assert_allclose', 'assert_array_almost_equal_nulp',
39    'assert_raises'
40    ]
41
42__all__ = __all__masked + __some__from_testing  # noqa: PLE0605
43
44
45def approx(a, b, fill_value=True, rtol=1e-5, atol=1e-8):
46    """
47    Returns true if all components of a and b are equal to given tolerances.
48
49    If fill_value is True, masked values considered equal. Otherwise,
50    masked values are considered unequal.  The relative error rtol should
51    be positive and << 1.0 The absolute error atol comes into play for
52    those elements of b that are very small or zero; it says how small a
53    must be also.
54
55    """
56    m = mask_or(getmask(a), getmask(b))
57    d1 = filled(a)
58    d2 = filled(b)
59    if d1.dtype.char == "O" or d2.dtype.char == "O":
60        return np.equal(d1, d2).ravel()
61    x = filled(
62        masked_array(d1, copy=False, mask=m), fill_value
63    ).astype(np.float64)
64    y = filled(masked_array(d2, copy=False, mask=m), 1).astype(np.float64)
65    d = np.less_equal(umath.absolute(x - y), atol + rtol * umath.absolute(y))
66    return d.ravel()
67
68
69def almost(a, b, decimal=6, fill_value=True):
70    """
71    Returns True if a and b are equal up to decimal places.
72
73    If fill_value is True, masked values considered equal. Otherwise,
74    masked values are considered unequal.
75
76    """
77    m = mask_or(getmask(a), getmask(b))
78    d1 = filled(a)
79    d2 = filled(b)
80    if d1.dtype.char == "O" or d2.dtype.char == "O":
81        return np.equal(d1, d2).ravel()
82    x = filled(
83        masked_array(d1, copy=False, mask=m), fill_value
84    ).astype(np.float64)
85    y = filled(masked_array(d2, copy=False, mask=m), 1).astype(np.float64)
86    d = np.around(np.abs(x - y), decimal) <= 10.0 ** (-decimal)
87    return d.ravel()
88
89
90def _assert_equal_on_sequences(actual, desired, err_msg=''):
91    """
92    Asserts the equality of two non-array sequences.
93
94    """
95    assert_equal(len(actual), len(desired), err_msg)
96    for k in range(len(desired)):
97        assert_equal(actual[k], desired[k], f'item={k!r}\n{err_msg}')
98
99
100def assert_equal_records(a, b):
101    """
102    Asserts that two records are equal.
103
104    Pretty crude for now.
105
106    """
107    assert_equal(a.dtype, b.dtype)
108    for f in a.dtype.names:
109        (af, bf) = (operator.getitem(a, f), operator.getitem(b, f))
110        if not (af is masked) and not (bf is masked):
111            assert_equal(operator.getitem(a, f), operator.getitem(b, f))
112
113
114def assert_equal(actual, desired, err_msg=''):
115    """
116    Asserts that two items are equal.
117
118    """
119    # Case #1: dictionary .....
120    if isinstance(desired, dict):
121        if not isinstance(actual, dict):
122            raise AssertionError(repr(type(actual)))
123        assert_equal(len(actual), len(desired), err_msg)
124        for k in desired:
125            if k not in actual:
126                raise AssertionError(f"{k} not in {actual}")
127            assert_equal(actual[k], desired[k], f'key={k!r}\n{err_msg}')
128        return
129    # Case #2: lists .....
130    if isinstance(desired, (list, tuple)) and isinstance(actual, (list, tuple)):
131        return _assert_equal_on_sequences(actual, desired, err_msg='')
132    if not (isinstance(actual, ndarray) or isinstance(desired, ndarray)):
133        msg = build_err_msg([actual, desired], err_msg,)
134        if not desired == actual:
135            raise AssertionError(msg)
136        return
137    # Case #4. arrays or equivalent
138    if ((actual is masked) and not (desired is masked)) or \
139            ((desired is masked) and not (actual is masked)):
140        msg = build_err_msg([actual, desired],
141                            err_msg, header='', names=('x', 'y'))
142        raise ValueError(msg)
143    actual = np.asanyarray(actual)
144    desired = np.asanyarray(desired)
145    (actual_dtype, desired_dtype) = (actual.dtype, desired.dtype)
146    if actual_dtype.char == "S" and desired_dtype.char == "S":
147        return _assert_equal_on_sequences(actual.tolist(),
148                                          desired.tolist(),
149                                          err_msg='')
150    return assert_array_equal(actual, desired, err_msg)
151
152
153def fail_if_equal(actual, desired, err_msg='',):
154    """
155    Raises an assertion error if two items are equal.
156
157    """
158    if isinstance(desired, dict):
159        if not isinstance(actual, dict):
160            raise AssertionError(repr(type(actual)))
161        fail_if_equal(len(actual), len(desired), err_msg)
162        for k in desired:
163            if k not in actual:
164                raise AssertionError(repr(k))
165            fail_if_equal(actual[k], desired[k], f'key={k!r}\n{err_msg}')
166        return
167    if isinstance(desired, (list, tuple)) and isinstance(actual, (list, tuple)):
168        fail_if_equal(len(actual), len(desired), err_msg)
169        for k in range(len(desired)):
170            fail_if_equal(actual[k], desired[k], f'item={k!r}\n{err_msg}')
171        return
172    if isinstance(actual, np.ndarray) or isinstance(desired, np.ndarray):
173        return fail_if_array_equal(actual, desired, err_msg)
174    msg = build_err_msg([actual, desired], err_msg)
175    if not desired != actual:
176        raise AssertionError(msg)
177
178
179assert_not_equal = fail_if_equal
180
181
182def assert_almost_equal(actual, desired, decimal=7, err_msg='', verbose=True):
183    """
184    Asserts that two items are almost equal.
185
186    The test is equivalent to abs(desired-actual) < 0.5 * 10**(-decimal).
187
188    """
189    if isinstance(actual, np.ndarray) or isinstance(desired, np.ndarray):
190        return assert_array_almost_equal(actual, desired, decimal=decimal,
191                                         err_msg=err_msg, verbose=verbose)
192    msg = build_err_msg([actual, desired],
193                        err_msg=err_msg, verbose=verbose)
194    if not round(abs(desired - actual), decimal) == 0:
195        raise AssertionError(msg)
196
197
198assert_close = assert_almost_equal
199
200
201def assert_array_compare(comparison, x, y, err_msg='', verbose=True, header='',
202                         fill_value=True):
203    """
204    Asserts that comparison between two masked arrays is satisfied.
205
206    The comparison is elementwise.
207
208    """
209    # Allocate a common mask and refill
210    m = mask_or(getmask(x), getmask(y))
211    x = masked_array(x, copy=False, mask=m, keep_mask=False, subok=False)
212    y = masked_array(y, copy=False, mask=m, keep_mask=False, subok=False)
213    if ((x is masked) and not (y is masked)) or \
214            ((y is masked) and not (x is masked)):
215        msg = build_err_msg([x, y], err_msg=err_msg, verbose=verbose,
216                            header=header, names=('x', 'y'))
217        raise ValueError(msg)
218    # OK, now run the basic tests on filled versions
219    return np.testing.assert_array_compare(comparison,
220                                           x.filled(fill_value),
221                                           y.filled(fill_value),
222                                           err_msg=err_msg,
223                                           verbose=verbose, header=header)
224
225
226def assert_array_equal(x, y, err_msg='', verbose=True):
227    """
228    Checks the elementwise equality of two masked arrays.
229
230    """
231    assert_array_compare(operator.__eq__, x, y,
232                         err_msg=err_msg, verbose=verbose,
233                         header='Arrays are not equal')
234
235
236def fail_if_array_equal(x, y, err_msg='', verbose=True):
237    """
238    Raises an assertion error if two masked arrays are not equal elementwise.
239
240    """
241    def compare(x, y):
242        return (not np.all(approx(x, y)))
243    assert_array_compare(compare, x, y, err_msg=err_msg, verbose=verbose,
244                         header='Arrays are not equal')
245
246
247def assert_array_approx_equal(x, y, decimal=6, err_msg='', verbose=True):
248    """
249    Checks the equality of two masked arrays, up to given number odecimals.
250
251    The equality is checked elementwise.
252
253    """
254    def compare(x, y):
255        "Returns the result of the loose comparison between x and y)."
256        return approx(x, y, rtol=10. ** -decimal)
257    assert_array_compare(compare, x, y, err_msg=err_msg, verbose=verbose,
258                         header='Arrays are not almost equal')
259
260
261def assert_array_almost_equal(x, y, decimal=6, err_msg='', verbose=True):
262    """
263    Checks the equality of two masked arrays, up to given number odecimals.
264
265    The equality is checked elementwise.
266
267    """
268    def compare(x, y):
269        "Returns the result of the loose comparison between x and y)."
270        return almost(x, y, decimal)
271    assert_array_compare(compare, x, y, err_msg=err_msg, verbose=verbose,
272                         header='Arrays are not almost equal')
273
274
275def assert_array_less(x, y, err_msg='', verbose=True):
276    """
277    Checks that x is smaller than y elementwise.
278
279    """
280    assert_array_compare(operator.__lt__, x, y,
281                         err_msg=err_msg, verbose=verbose,
282                         header='Arrays are not less-ordered')
283
284
285def assert_mask_equal(m1, m2, err_msg=''):
286    """
287    Asserts the equality of two masks.
288
289    """
290    if m1 is nomask:
291        assert_(m2 is nomask)
292    if m2 is nomask:
293        assert_(m1 is nomask)
294    assert_array_equal(m1, m2, err_msg=err_msg)
295