1# -*- coding: utf-8 -*-
2
3from ..regression import least_squares, irls
4from ..testing import requires
5
6
7@requires("numpy")
8def test_irls():
9    import numpy as np
10
11    x = np.linspace(0, 100)
12    y = np.exp(-x / 47)
13    b, c, info = irls(x, np.log(y))
14    assert abs(b[1] + 1 / 47) < 1e-5
15    assert np.all(c < 1e-4)
16    assert info["success"] is True
17    assert info["niter"] < 3
18
19
20@requires("numpy")
21def test_least_squares():
22    import numpy as np
23
24    x, y, w = [0, 1, 2], [0, 1, 2], [1, 1, 1]
25    beta, vcv, r2 = least_squares(x, y)
26    assert np.allclose(beta, [0, 1])
27    assert np.allclose(vcv, 0)
28    assert np.allclose(r2, 1)
29
30    wbeta, wvcv, wr2 = least_squares(x, y, w)
31    assert np.allclose(wbeta, [0, 1])
32    assert np.allclose(wvcv, 0)
33    assert np.allclose(wr2, 1)
34