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

Code before:
        p = BarycentricInterpolator(x, y)
        assert_allclose(r(xx), p(xx), rtol=1e-12, atol=1e-12)

    @pytest.mark.parametrize("y_shape", [(2,), (2, 3, 1), (1, 5, 6, 4)])
    @pytest.mark.parametrize("xx_shape", [(100), (10, 10)])
    def test_trailing_dim(self, y_shape, xx_shape):
        x = np.linspace(0, 1)
        y = np.broadcast_to(
            np.expand_dims(np.sin(x), tuple(range(1, len(y_shape) + 1))),
            x.shape + y_shape
        )

        r = FloaterHormannInterpolator(x, y)

        rng = np.random.default_rng(897138947238097528091759187597)
        xx = rng.random(xx_shape)
        yy = np.broadcast_to(
            np.expand_dims(np.sin(xx), tuple(range(xx.ndim, len(y_shape) + xx.ndim))),
            xx.shape + y_shape
        )
        rr = r(xx)
        assert rr.shape == xx.shape + y_shape
        assert_allclose(rr, yy, rtol=1e-6)

        pytest.raises(NotImplementedError, r.roots)
        pytest.raises(NotImplementedError, r.residues)

    def test_zeros(self):
        x = np.linspace(0, 10, num=100)
        r = FloaterHormannInterpolator(x, np.sin(np.pi*x))

        err = np.abs(np.subtract.outer(r.roots(), np.arange(11))).min(axis=0)
        assert_array_less(err, 1e-5)

    def test_no_poles(self):
        x = np.linspace(-1, 1)
        r = FloaterHormannInterpolator(x, 1/x**2)
        p = r.poles()
        mask = (p.real >= -1) & (p.real <= 1) & (np.abs(p.imag) < 1.e-12)
        assert np.sum(mask) == 0

    @pytest.mark.parametrize('eval_shape', [(), (1,), (3,)])
    @pytest.mark.parametrize('axis', [-1, 0, 1])
    def test_batch(self, eval_shape, axis):
        rng = np.random.default_rng(4329872134985134)
        n = 10
        shape = (2, 3, 4, n)
        domain = (0, 10)

        x = np.linspace(*domain, n)
        y = np.moveaxis(rng.random(shape), -1, axis)

Review comment (performance):
`poles`, too, for completeness?
And why add these here instead of the new batch test?

Code after:
        zz = xx + xx*1j
        assert_allclose(r(zz), np.sin(zz), rtol=1e-12)

    def test_polyinterp(self):
        # check that when d=n-1 FH gives a polynomial interpolant
        x = np.linspace(0, 1, 11)
        xx = np.linspace(0, 1, 1001)
        y = np.sin(x)
        r = FloaterHormannInterpolator(x, y, d=x.size-1)
        p = BarycentricInterpolator(x, y)
        assert_allclose(r(xx), p(xx), rtol=1e-12, atol=1e-12)

    @pytest.mark.parametrize("y_shape", [(2,), (2, 3, 1), (1, 5, 6, 4)])
    @pytest.mark.parametrize("xx_shape", [(100), (10, 10)])
    def test_trailing_dim(self, y_shape, xx_shape):
        x = np.linspace(0, 1)
        y = np.broadcast_to(
            np.expand_dims(np.sin(x), tuple(range(1, len(y_shape) + 1))),
            x.shape + y_shape
        )

        r = FloaterHormannInterpolator(x, y)

        rng = np.random.default_rng(897138947238097528091759187597)
        xx = rng.random(xx_shape)
        yy = np.broadcast_to(
            np.expand_dims(np.sin(xx), tuple(range(xx.ndim, len(y_shape) + xx.ndim))),
            xx.shape + y_shape
        )
        rr = r(xx)
        assert rr.shape == xx.shape + y_shape
        assert_allclose(rr, yy, rtol=1e-6)


    def test_zeros(self):
        x = np.linspace(0, 10, num=100)
        r = FloaterHormannInterpolator(x, np.sin(np.pi*x))

        err = np.abs(np.subtract.outer(r.roots(), np.arange(11))).min(axis=0)
        assert_array_less(err, 1e-5)

    def test_no_poles(self):
        x = np.linspace(-1, 1)
        r = FloaterHormannInterpolator(x, 1/x**2)
        p = r.poles()
        mask = (p.real >= -1) & (p.real <= 1) & (np.abs(p.imag) < 1.e-12)
        assert np.sum(mask) == 0

    @pytest.mark.parametrize('eval_shape', [(), (1,), (3,)])
    @pytest.mark.parametrize('axis', [-1, 0, 1])
    def test_batch(self, eval_shape, axis):