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

Code before:
        np.testing.assert_allclose(p_jac[:, -1],
            np.where(n >= 1, -pow(x, n) * np.inf, 0))

        for m in range(3, m_max + 1):
            np.testing.assert_allclose(p_jac[:, m], 0)
            np.testing.assert_allclose(p_jac[:, -m], 0)

    def test_assoc_legendre_p_norm_m0(self):
        # gh 24099
        m = 0
        n_max = 10
        z = np.linspace(-1, 1, 101)
        for n in range(n_max + 1):
            p = assoc_legendre_p(n, m, z, norm=False)
            p_norm = assoc_legendre_p(n, m, z, norm=True)
            factor = np.sqrt((2 * n + 1) / 2)
            np.testing.assert_allclose(p_norm, factor * p)

    def test_assoc_legendre_p_all_norm_m0(self):
        # gh 24099
        m = 0
        n_max = 10
        z = np.linspace(-1, 1, 101)
        p_all = assoc_legendre_p_all(n_max, m, z, norm=False)[0, :, 0, :]
        p_norm_all = assoc_legendre_p_all(n_max, m, z, norm=True)[0, :, 0, :]
        for n in range(n_max + 1):
            factor = np.sqrt((2 * n + 1) / 2)
            np.testing.assert_allclose(p_norm_all[n, :], factor * p_all[n, :])


class TestMultiAssocLegendreP:
    @pytest.mark.parametrize("shape", [(1000,), (4, 9), (3, 5, 7)])
    @pytest.mark.parametrize("branch_cut", [2, 3])
    @pytest.mark.parametrize("z_min, z_max", [(-10 - 10j, 10 + 10j),
        (-1, 1), (-10j, 10j)])
    @pytest.mark.parametrize("norm", [True, False])
    def test_specific(self, shape, branch_cut, z_min, z_max, norm):
        rng = np.random.default_rng(1234)

        z = rng.uniform(z_min.real, z_max.real, shape) + \
            1j * rng.uniform(z_min.imag, z_max.imag, shape)

        p, p_jac = assoc_legendre_p_all(4, 4,
            z, branch_cut=branch_cut, norm=norm, diff_n=1)

        np.testing.assert_allclose(p[0, 0],
            assoc_legendre_p_0_0(z, branch_cut=branch_cut, norm=norm))
        np.testing.assert_allclose(p[0, 1], 0)
        np.testing.assert_allclose(p[0, 2], 0)
        np.testing.assert_allclose(p[0, 3], 0)
        np.testing.assert_allclose(p[0, 4], 0)

Review comment (bug):
You could parametrize both tests over `n` to get better testing feedback if only a subset of cases start failing at some point, but not sure that is even worth flushing the CI again over.

Code after:
        np.testing.assert_allclose(p_jac[:, -1],
            np.where(n >= 1, -pow(x, n) * np.inf, 0))

        for m in range(3, m_max + 1):
            np.testing.assert_allclose(p_jac[:, m], 0)
            np.testing.assert_allclose(p_jac[:, -m], 0)

    @pytest.mark.parametrize("n", np.arange(0, 20))
    def test_assoc_legendre_p_norm_m0(self, n):
        # gh 24099
        m = 0
        z = np.linspace(-1, 1, 101)
        p = assoc_legendre_p(n, m, z, norm=False)
        p_norm = assoc_legendre_p(n, m, z, norm=True)
        factor = np.sqrt((2 * n + 1) / 2)
        np.testing.assert_allclose(p_norm, factor * p)

    @pytest.mark.parametrize("n", np.arange(0, 20))
    def test_assoc_legendre_p_all_norm_m0(self, n):
        # gh 24099
        m = 0
        z = np.linspace(-1, 1, 101)
        p_all = assoc_legendre_p_all(n, m, z, norm=False)[0, : n + 1, 0, :]
        p_norm_all = assoc_legendre_p_all(n, m, z, norm=True)[0, : n + 1, 0, :]
        factor = np.sqrt((2 * np.arange(n + 1) + 1) / 2)
        factor = factor[:, np.newaxis]
        np.testing.assert_allclose(p_norm_all, factor * p_all)


class TestMultiAssocLegendreP:
    @pytest.mark.parametrize("shape", [(1000,), (4, 9), (3, 5, 7)])
    @pytest.mark.parametrize("branch_cut", [2, 3])
    @pytest.mark.parametrize("z_min, z_max", [(-10 - 10j, 10 + 10j),
        (-1, 1), (-10j, 10j)])
    @pytest.mark.parametrize("norm", [True, False])
    def test_specific(self, shape, branch_cut, z_min, z_max, norm):
        rng = np.random.default_rng(1234)

        z = rng.uniform(z_min.real, z_max.real, shape) + \
            1j * rng.uniform(z_min.imag, z_max.imag, shape)

        p, p_jac = assoc_legendre_p_all(4, 4,
            z, branch_cut=branch_cut, norm=norm, diff_n=1)

        np.testing.assert_allclose(p[0, 0],
            assoc_legendre_p_0_0(z, branch_cut=branch_cut, norm=norm))
        np.testing.assert_allclose(p[0, 1], 0)
        np.testing.assert_allclose(p[0, 2], 0)
        np.testing.assert_allclose(p[0, 3], 0)
        np.testing.assert_allclose(p[0, 4], 0)
        np.testing.assert_allclose(p[0, -4], 0)