Repository: scipy/scipy
File: scipy/interpolate/tests/test_bary_rational.py (Python)

Code before:
    def test_spiral_cleanup(self):
        z = np.exp(np.linspace(-0.5, 0.5 + 15j*np.pi, num=1000))
        # here we set `rtol=0` to force froissart doublets, without cleanup there
        # are many spurious poles
        with pytest.warns(RuntimeWarning):
            r = AAA(z, np.tan(np.pi*z/2), rtol=0, max_terms=60, clean_up=False)
        n_spurious = np.sum(np.abs(r.residues()) < 1e-14)
        with pytest.warns(RuntimeWarning):
            assert r.clean_up() >= 1
        # check there are less potentially spurious poles than before
        assert np.sum(np.abs(r.residues()) < 1e-14) < n_spurious
        # check accuracy
        assert_allclose(r(z), np.tan(np.pi*z/2), atol=6e-12, rtol=3e-12)


class BatchFloaterHormann:
    # BSpline-line class with reference batch behavior
    def __init__(self, x, y, axis):
        y = np.moveaxis(y, axis, -1)
        self._batch_shape = y.shape[:-1]
        self._interps = [FloaterHormannInterpolator(x, yi,)
                         for yi in y.reshape(-1, y.shape[-1])]
        self._axis = axis

    def __call__(self, x):
        y = [intep(x) for intep in self._interps]
        y = np.reshape(y, self._batch_shape + x.shape)
        return np.moveaxis(y, -1, self._axis) if x.shape else y


class TestFloaterHormann:
    def runge(self, z):
        return 1/(1 + z**2)

    def scale(self, n, d):
        return (-1)**(np.arange(n) + d) * factorial(d)

    def test_iv(self):
        with pytest.raises(ValueError, match="`x`"):
            FloaterHormannInterpolator([[0]], [0], d=0)
        with pytest.raises(ValueError, match="`y`"):
            FloaterHormannInterpolator([0], 0, d=0)
        with pytest.raises(ValueError, match="size"):
            FloaterHormannInterpolator([0], [[1, 1], [1, 1]], d=0)
        with pytest.raises(ValueError, match="finite"):
            FloaterHormannInterpolator([np.inf], [1], d=0)
        with pytest.raises(ValueError, match="`d`"):
            FloaterHormannInterpolator([0], [0], d=-1)
        with pytest.raises(ValueError, match="`d`"):
            FloaterHormannInterpolator([0], [0], d=10)
        with pytest.raises(TypeError):

Review comment (suggestion):
```suggestion
        y = [interp(x) for interp in self._interps]
```

Code after:
    def test_spiral_cleanup(self):
        z = np.exp(np.linspace(-0.5, 0.5 + 15j*np.pi, num=1000))
        # here we set `rtol=0` to force froissart doublets, without cleanup there
        # are many spurious poles
        with pytest.warns(RuntimeWarning):
            r = AAA(z, np.tan(np.pi*z/2), rtol=0, max_terms=60, clean_up=False)
        n_spurious = np.sum(np.abs(r.residues()) < 1e-14)
        with pytest.warns(RuntimeWarning):
            assert r.clean_up() >= 1
        # check there are less potentially spurious poles than before
        assert np.sum(np.abs(r.residues()) < 1e-14) < n_spurious
        # check accuracy
        assert_allclose(r(z), np.tan(np.pi*z/2), atol=6e-12, rtol=3e-12)

    def test_diag_scaling(self):
        # fails without diag scaling
        z = np.logspace(-15, 0, 300)
        f = np.sqrt(z)
        r = AAA(z, f)

        zz = np.logspace(-15, 0, 500)
        assert_allclose(r(zz), np.sqrt(zz), rtol=9e-6)


class BatchFloaterHormann:
    # FloaterHormann class with reference batch behaviour
    def __init__(self, x, y, axis):
        y = np.moveaxis(y, axis, -1)
        self._batch_shape = y.shape[:-1]
        self._interps = [FloaterHormannInterpolator(x, yi,)
                         for yi in y.reshape(-1, y.shape[-1])]
        self._axis = axis

    def __call__(self, x):
        y = [interp(x) for interp in self._interps]
        y = np.reshape(y, self._batch_shape + x.shape)
        return np.moveaxis(y, -1, self._axis) if x.shape else y


class TestFloaterHormann:
    def runge(self, z):
        return 1/(1 + z**2)

    def scale(self, n, d):
        return (-1)**(np.arange(n) + d) * factorial(d)

    def test_iv(self):
        with pytest.raises(ValueError, match="`x`"):
            FloaterHormannInterpolator([[0]], [0], d=0)
        with pytest.raises(ValueError, match="`y`"):
            FloaterHormannInterpolator([0], 0, d=0)