Repository: scipy/scipy
File: scipy/special/tests/test_legendre.py (Python)

Code before:

    @pytest.mark.parametrize("n_max", [1, 2, 4, 8, 16, 32])
    @pytest.mark.parametrize("x_shape", [(10,), (4, 9), (3, 5, 7)])
    def test_all_ode(self, n_max, x_shape):
        rng = np.random.default_rng(1234)

        x = rng.uniform(-1, 1, x_shape)
        p, p_jac, p_hess = legendre_p_all(n_max, x, diff_n=2)

        n = np.arange(n_max + 1)
        n = np.expand_dims(n, axis = tuple(range(1, x.ndim + 1)))

        assert p.shape == (len(n),) + x.shape
        assert p_jac.shape == p.shape
        assert p_hess.shape == p_jac.shape

        err = (1 - x * x) * p_hess - 2 * x * p_jac + n * (n + 1) * p
        np.testing.assert_allclose(err, 0, atol=1e-10)


class TestAssocLegendreP:
    def test_assoc_legendre_bug(self):
        """
        This test detects the bug reported in
        https://github.com/scipy/scipy/issues/23101
        """
        z = np.array([-1, -.5, 0, .5, 1])
        expected = assoc_legendre_p_1_0(z)
        result = assoc_legendre_p(1, 0, z)
        assert_allclose(np.squeeze(result), expected)

        expected = assoc_legendre_p_3_0(z)
        result = assoc_legendre_p(3, 0, z)
        assert_allclose(np.squeeze(result), expected)

    @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7, 10)])
    @pytest.mark.parametrize("m_max", [5, 4])
    @pytest.mark.parametrize("n_max", [7, 10])
    def test_lpmn(self, shape, n_max, m_max):
        rng = np.random.default_rng(1234)

        x = rng.uniform(-0.99, 0.99, shape)
        p_all, p_all_jac, p_all_hess = \
            assoc_legendre_p_all(n_max, m_max, x, diff_n=2)

        n = np.arange(n_max + 1)
        n = np.expand_dims(n, axis = tuple(range(1, x.ndim + 2)))

        m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)])
        m = np.expand_dims(m, axis = (0,) + tuple(range(2, x.ndim + 2)))


Review comment (suggestion):
this is more concise and makes the test name more descriptive
```suggestion
    def test_assoc_legendre_gh23101(self):
```

Code after:

    @pytest.mark.parametrize("n_max", [1, 2, 4, 8, 16, 32])
    @pytest.mark.parametrize("x_shape", [(10,), (4, 9), (3, 5, 7)])
    def test_all_ode(self, n_max, x_shape):
        rng = np.random.default_rng(1234)

        x = rng.uniform(-1, 1, x_shape)
        p, p_jac, p_hess = legendre_p_all(n_max, x, diff_n=2)

        n = np.arange(n_max + 1)
        n = np.expand_dims(n, axis = tuple(range(1, x.ndim + 1)))

        assert p.shape == (len(n),) + x.shape
        assert p_jac.shape == p.shape
        assert p_hess.shape == p_jac.shape

        err = (1 - x * x) * p_hess - 2 * x * p_jac + n * (n + 1) * p
        np.testing.assert_allclose(err, 0, atol=1e-10)


class TestAssocLegendreP:
    def test_assoc_legendre_gh23101(self):
        z = np.array([-1, -.5, 0, .5, 1])
        expected = assoc_legendre_p_1_0(z)
        result = assoc_legendre_p(1, 0, z)
        assert_allclose(np.squeeze(result), expected)

        expected = assoc_legendre_p_3_0(z)
        result = assoc_legendre_p(3, 0, z)
        assert_allclose(np.squeeze(result), expected)

    @pytest.mark.parametrize("shape", [(10,), (4, 9), (3, 5, 7, 10)])
    @pytest.mark.parametrize("m_max", [5, 4])
    @pytest.mark.parametrize("n_max", [7, 10])
    def test_lpmn(self, shape, n_max, m_max):
        rng = np.random.default_rng(1234)

        x = rng.uniform(-0.99, 0.99, shape)
        p_all, p_all_jac, p_all_hess = \
            assoc_legendre_p_all(n_max, m_max, x, diff_n=2)

        n = np.arange(n_max + 1)
        n = np.expand_dims(n, axis = tuple(range(1, x.ndim + 2)))

        m = np.concatenate([np.arange(m_max + 1), np.arange(-m_max, 0)])
        m = np.expand_dims(m, axis = (0,) + tuple(range(2, x.ndim + 2)))

        x = np.expand_dims(x, axis = (0, 1))
        p, p_jac, p_hess = assoc_legendre_p(n, m, x, diff_n=2)

        np.testing.assert_allclose(p, p_all)