Repository: fatiando/harmonica
Pull request #505 (2024-05-10 15:33:42)

Add associated Legendre function calculations
The numba-based functions calculate the associated Legendre functions (unnormalized, Schmidt normalized, and fully normalized) and their derivatives $\dfrac{\partial P_n^m}{\partial \theta}(\cos\theta)$. The derivatives functions can be used to calculate higher order derivatives as well. Values are tested against analytical solutions for the first 4 degrees as well as an identity for the Schmidt functions and the Legendre equation for the Schmidt and fully normalized functions. The values are returned as 2D numpy arrays. This is a bit of a waste of memory but it's much easier to use than storing these in sparse matrices or 1D arrays. But it should be acceptable since we're unlikely to ever go over ~2700 degree so the waste is in the order of a few Mb at worse.

<!--
Thank you for contributing a pull request to Fatiando! 💖

👆🏽 ABOVE: Describe the changes proposed and WHY you made them.

👇🏽 BELOW: Link to any relevant issue or pull request.

Please ensure you have taken a look at the CONTRIBUTING.md file 
in this repository (if available) and the general guidelines at 
https://github.com/fatiando/community/blob/main/CONTRIBUTING.md
-->

**Relevant issues/PRs:**
<!--
Example: "Fixes #1234" / "See also #345" / "Relevant to #111"
Use keywords (e.g., Fixes, Closes) to create the links and automatically
close issues when this PR is merged. 
See https://github.com/blog/1506-closing-issues-via-pull-requests
-->
Related to #504 

leouieda: @santisoler and @MarkWieczorek this is ready now. I added the Holmes and Featherstone scaling and tested until degree 2800. I was surprised that calculating the derivatives on rescaled functions works well. That's nice because it means we can keep them separate and use the derivative functions to calculate higher order derivatives (which would appear in gravity gradients).
leouieda: Thanks, @santisoler! What do you think?

Patch:
diff --git a/harmonica/_spherical_harmonics/legendre.py b/harmonica/_spherical_harmonics/legendre.py
new file mode 100644
index 0000000..7e25a3d
--- /dev/null
+++ b/harmonica/_spherical_harmonics/legendre.py
@@ -0,0 +1,383 @@
+# Copyright (c) 2018 The Harmonica Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+"""
+Calculation of associated Legendre functions and their derivatives.
+"""
+import numba
+import numpy as np
+
+
+@numba.jit(nopython=True)
+def _rescale(u, max_degree, p):
+    "Rescale Legendre functions to their original range"
+    rescale = 1e280
+    for m in range(0, max_degree + 1):
+        if m > 0:
+            rescale *= u
+        for n in range(m, max_degree + 1):
+            p[n, m] *= rescale
+
+
+@numba.jit(nopython=True)
+def associated_legendre(x, max_degree, p):
+    """
+    Unnormalized associated Legendre functions up to a maximum degree.
+
+    The functions :math:`P_n^m` are solutions to Legendre's equation. We
+    calculate their values using the recursive relations defined in Alken
+    (2022).
+
+    .. note::
+
+        This function does not include the Condon-Shortly phase.
+
+    .. warning::
+
+        Unnormalized functions should not be used beyond degree 13 because they
+        will overflow 64-bit float range.
+
+    Parameters
+    ----------
+    x : float
+        The argument of :math:`P_n^m(x)`. Must be in the range [-1, 1].
+    max_degree : int
+        The maximum degree for the calculation.
+    p : numpy.array
+        A 2D array with shape ``(max_degree + 1, max_degree + 1)`` that will be
+        filled with the output values.
+
+    References
+    ----------
+
+    Alken, Patrick (2022). GSL Technical Report #1 - GSL-TR-001-20220827 -
+      Implementation of associated Legendre functions in GSL.
+      https://www.gnu.org/software/gsl/tr/tr001.pdf
+    """
+    u = np.sqrt((1 - x) * (1 + x))
+    p[0, 0] = 1
+    for n in range(1, max_degree + 1):
+        for m in range(0, n - 1):
+            a_nm = (2 * n - 1) / (n - m)
+            b_nm = -(n + m - 1) / (n - m)
+            p[n, m] = a_nm * x * p[n - 1, m] + b_nm * p[n - 2, m]
+        c_nm = 2 * n - 1
+        p[n, n - 1] = c_nm * x * p[n - 1, n - 1]
+        d_nm = 2 * n - 1
+        p[n, n] = d_nm * u * p[n - 1, n - 1]
+
+
+@numba.jit(nopython=True)
+def associated_legendre_derivative(max_degree, p, dp):
+    """
+    Derivatives in theta of unnormalized associated Legendre functions.
+
+    Calculates the derivative:
+
+    .. math::
+
+        \\dfrac{\\partial P_n^m}{\\partial \\theta}(\\cos \\theta)
+
+    using the recursive relations defined in Alken (2022).
+
+    Higher-order derivatives can be calculated by passing the output of this
+    function as the ``p`` argument.
+
+    .. note::
+
+        This function does not include the Condon-Shortly phase.
+
+    .. warning::
+
+        Unnormalized functions should not be used beyond degree 13 because they
+        will overflow 64-bit float range.
+
+    Parameters
+    ----------
+    max_degree : int
+        The maximum degree for the calculation.
+    p : numpy.ndarray
+        A 2D array with the unnormalized associated Legendre functions
+        calculated for :math:`\\cos\\theta`.
+    dp : numpy.array
+        A 2D array with shape ``(max_degree + 1, max_degree + 1)`` that will be
+        filled with the output values.
+
+    References
+    ----------
+
+    Alken, Patrick (2022). GSL Technical Report #1 - GSL-TR-001-20220827 -
+      Implementation of associated Legendre functions in GSL.
+      https://www.gnu.org/software/gsl/tr/tr001.pdf
+    """
+    dp[0, 0] = 0
+    for n in range(1, max_degree + 1):
+        a_nm = -1
+        dp[n, 0] = a_nm * p[n, 1]
+        for m in range(1, n):
+            b_nm = 0.5 * (n + m) * (n - m + 1)
+            c_nm = -0.5
+            dp[n, m] = b_nm * p[n, m - 1] + c_nm * p[n, m + 1]
+        d_nm = n
+        dp[n, n] = d_nm * p[n, n - 1]
+
+
+@numba.jit(nopython=True)
+def associated_legendre_schmidt(x, max_degree, p):
+    """
+    Schmidt normalized associated Legendre functions up to maximum degree.
+
+    The functions :math:`P_n^m` are solutions to Legendre's equation. The
+    Schmidt normalization is applied to the functions to constrain their value
+    range. This normalization if often used in geomagnetic field models. We
+    calculate their values using the recursive relations defined in Alken
+    (2022).
+
+    .. note::
+
+        This function does not include the Condon-Shortly phase.
+
+    .. note::
+
+        This function uses the scaling scheme of Holmes and Featherstone (2002)
+        and produces accurate results until degree and order 2800.
+
+    Parameters
+    ----------
+    x : float
+        The argument of :math:`P_n^m(x)`. Must be in the range [-1, 1].
+    max_degree : int
+        The maximum degree for the calculation.
+    p : numpy.array
+        A 2D array with shape ``(max_degree + 1, max_degree + 1)`` that will be
+        filled with the output values.
+
+    References
+    ----------
+
+    Alken, Patrick (2022). GSL Technical Report #1 - GSL-TR-001-20220827 -
+      Implementation of associated Legendre functions in GSL.
+      https://www.gnu.org/software/gsl/tr/tr001.pdf
+    """
+    u = np.sqrt((1 - x) * (1 + x))
+    # Pre-compute square roots of integers used in the loops
+    sqrt = np.sqrt(np.arange(2 * (max_degree + 1)))
+    # Use the Holmes and Featherstone (2002) scaling to compute scaled Pnm
+    # All terms are scaled by the max float range 1e280
+    p[0, 0] = 1 * 1e-280
+    p[1, 0] = x * 1e-280
+    # This equation doesn't have u because of the scaling
+    p[1, 1] = p[0, 0]
+    # Calculate the zero order terms first
+    for n in range(2, max_degree + 1):
+        a_n0 = (2 * n - 1) / (sqrt[n] ** 2)
+        b_n0 = -((sqrt[n - 1] / sqrt[n]) ** 2)
+        p[n, 0] = a_n0 * x * p[n - 1, 0] + b_n0 * p[n - 2, 0]
+    # Now calculate the other terms
+    for n in range(2, max_degree + 1):
+        for m in range(1, n - 1):
+            a_nm = (2 * n - 1) / (sqrt[n + m] * sqrt[n - m])
+            b_nm = -sqrt[n + m - 1] * sqrt[n - m - 1] / (sqrt[n + m] * sqrt[n - m])
+            p[n, m] = a_nm * x * p[n - 1, m] + b_nm * p[n - 2, m]
+        c_nm = sqrt[2 * n - 1]
+        p[n, n - 1] = c_nm * x * p[n - 1, n - 1]
+        d_nm = sqrt[2 * n - 1] / sqrt[2 * n]
+        # This equation doesn't have u because of the scaling
+        p[n, n] = d_nm * p[n - 1, n - 1]
+    # Now return everything to the original float range and rescale by u**m
+    _rescale(u, max_degree, p)
+
+
+@numba.jit(nopython=True)
+def associated_legendre_schmidt_derivative(max_degree, p, dp):
+    """
+    Derivatives in theta of Schmidt normalized associated Legendre functions.
+
+    Calculates the derivative:
+
+    .. math::
+
+        \\dfrac{\\partial P_n^m}{\\partial \\theta}(\\cos \\theta)
+
+    using the recursive relations defined in Alken (2022).
+
+    Higher-order derivatives can be calculated by passing the output of this
+    function as the ``p`` argument.
+
+    .. note::
+
+        This function does not include the Condon-Shortly phase.
+
+    .. note::
+
+        First and second order derivatives are tested against the Legendre
+        differential equation and produce accurate results until degree and
+        order 2800.
+
+    Parameters
+    ----------
+    max_degree : int
+        The maximum degree for the calculation.
+    p : numpy.ndarray
+        A 2D array with the Schmidt normalized associated Legendre functions
+        calculated for :math:`\\cos\\theta`.
+    dp : numpy.array
+        An 2D array with shape ``(max_degree + 1, max_degree + 1)`` that will
+        be filled with the output values.
+
+    References
+    ----------
+
+    Alken, Patrick (2022). GSL Technical Report #1 - GSL-TR-001-20220827 -
+      Implementation of associated Legendre functions in GSL.
+      https://www.gnu.org/software/gsl/tr/tr001.pdf
+    """
+    # Pre-compute square roots of integers used in the loops
+    sqrt = np.sqrt(np.arange(2 * (max_degree + 1)))
+    dp[0, 0] = 0
+    dp[1, 0] = -p[1, 1]
+    dp[1, 1] = p[1, 0]
+    for n in range(2, max_degree + 1):
+        a_nm = -sqrt[n] * sqrt[n + 1] / sqrt[2]
+        dp[n, 0] = a_nm * p[n, 1]
+        b_nm = 0.5 * sqrt[n] * sqrt[n + 1] * sqrt[2]
+        c_nm = -0.5 * sqrt[n + 2] * sqrt[n - 1]
+        dp[n, 1] = b_nm * p[n, 0] + c_nm * p[n, 2]
+        for m in range(2, n):
+            b_nm = 0.5 * sqrt[n + m] * sqrt[n - m + 1]
+            c_nm = -0.5 * sqrt[n + m + 1] * sqrt[n - m]
+            dp[n, m] = b_nm * p[n, m - 1] + c_nm * p[n, m + 1]
+        d_nm = 0.5 * sqrt[2] * sqrt[n]
+        dp[n, n] = d_nm * p[n, n - 1]
+
+
+@numba.jit(nopython=True)
+def associated_legendre_full(x, max_degree, p):
+    """
+    Fully normalized associated Legendre functions up to maximum degree.
+
+    The functions :math:`P_n^m` are solutions to Legendre's equation. The full
+    normalization is applied to the functions to constrain their value range.
+    This normalization if often used in gravity field models. We calculate
+    their values using the recursive relations defined in Alken (2022) but with
+    the definition of the normalization from Hofmann-Wellenhof and Moritz
+    (2006).
+
+    .. note::
+
+        This function does not include the Condon-Shortly phase.
+
+    Parameters
+    ----------
+    x : float
+        The argument of :math:`P_n^m(x)`. Must be in the range [-1, 1].
+    max_degree : int
+        The maximum degree for the calculation.
+    p : numpy.array
+        A 2D array with shape ``(max_degree + 1, max_degree + 1)`` that will be
+        filled with the output values.
+
+    References
+    ----------
+
+    Alken, Patrick (2022). GSL Technical Report #1 - GSL-TR-001-20220827 -
+      Implementation of associated Legendre functions in GSL.
+      https://www.gnu.org/software/gsl/tr/tr001.pdf
+    Hofmann-Wellenhof, B., & Moritz, H. (2006). Physical Geodesy (2nd, corr.
+      ed. 2006 edition ed.). Wien ; New York: Springer.
+    """
+    u = np.sqrt((1 - x) * (1 + x))
+    # Pre-compute square roots of integers used in the loops
+    sqrt = np.sqrt(np.arange(2 * (max_degree + 1)))
+    # Use the Holmes and Featherstone (2002) scaling to compute scaled Pnm
+    # All terms are scaled by the max float range 1e280
+    p[0, 0] = 1 * 1e-280
+    p[1, 0] = x * sqrt[3] * 1e-280
+    # This equation doesn't have u because of the scaling
+    p[1, 1] = p[0, 0] * sqrt[3]
+    # Calculate the zero order terms first
+    for n in range(2, max_degree + 1):
+        a_n0 = sqrt[2 * n - 1] * sqrt[2 * n + 1] / n
+        b_n0 = -(n - 1) * sqrt[2 * n + 1] / (n * sqrt[2 * n - 3])
+        p[n, 0] = a_n0 * x * p[n - 1, 0] + b_n0 * p[n - 2, 0]
+    # Now calculate the other terms
+    for n in range(2, max_degree + 1):
+        for m in range(1, n - 1):
+            a_nm = sqrt[2 * n + 1] * sqrt[2 * n - 1] / (sqrt[n + m] * sqrt[n - m])
+            b_nm = (
+                -sqrt[n + m - 1]
+                * sqrt[n - m - 1]
+                * sqrt[2 * n + 1]
+                / (sqrt[n + m] * sqrt[n - m] * sqrt[2 * n - 3])
+            )
+            p[n, m] = a_nm * x * p[n - 1, m] + b_nm * p[n - 2, m]
+        c_nm = sqrt[2 * n + 1]
+        p[n, n - 1] = c_nm * x * p[n - 1, n - 1]
+        d_nm = sqrt[2 * n + 1] / sqrt[2 * n]
+        # This equation doesn't have u because of the scaling
+        p[n, n] = d_nm * p[n - 1, n - 1]
+    # Now return everything to the original float range and rescale by u**m
+    _rescale(u, max_degree, p)
+
+
+@numba.jit(nopython=True)
+def associated_legendre_full_derivative(max_degree, p, dp):
+    """
+    Derivatives in theta of fully normalized associated Legendre functions.
+
+    Calculates the derivative:
+
+    .. math::
+
+        \\dfrac{\\partial P_n^m}{\\partial \\theta}(\\cos \\theta)
+
+    Uses the recursive relations defined in Alken (2022) but with the
+    definition of the normalization from Hofmann-Wellenhof and Moritz (2006).
+
+    Higher-order derivatives can be calculated by passing the output of this
+    function as the ``p`` argument.
+
+    .. note::
+
+        This function does not include the Condon-Shortly phase.
+
+    Parameters
+    ----------
+    max_degree : int
+        The maximum degree for the calculation.
+    p : numpy.ndarray
+        A 2D array with the fully normalized associated Legendre functions
+        calculated for :math:`\\cos\\theta`.
+    dp : numpy.array
+        A 2D array with shape ``(max_degree + 1, max_degree + 1)`` that will be
+        filled with the output values.
+
+    References
+    ----------
+
+    Alken, Patrick (2022). GSL Technical Report #1 - GSL-TR-001-20220827 -
+      Implementation of associated Legendre functions in GSL.
+      https://www.gnu.org/software/gsl/tr/tr001.pdf
+    Hofmann-Wellenhof, B., & Moritz, H. (2006). Physical Geodesy (2nd, corr.
+      ed. 2006 edition ed.). Wien ; New York: Springer.
+    """
+    # Pre-compute square roots of integers used in the loops
+    sqrt = np.sqrt(np.arange(2 * (max_degree + 1)))
+    dp[0, 0] = 0
+    dp[1, 0] = -p[1, 1]
+    dp[1, 1] = p[1, 0]
+    for n in range(2, max_degree + 1):
+        a_nm = -sqrt[n] * sqrt[n + 1] / sqrt[2]
+        dp[n, 0] = a_nm * p[n, 1]
+        b_n1 = 0.5 * sqrt[2] * sqrt[n + 1] * sqrt[n]
+        c_n1 = -0.5 * sqrt[n + 2] * sqrt[n - 1]
+        dp[n, 1] = b_n1 * p[n, 0] + c_n1 * p[n, 2]
+        for m in range(2, n):
+            b_nm = 0.5 * sqrt[n + m] * sqrt[n - m + 1]
+            c_nm = -0.5 * sqrt[n + m + 1] * sqrt[n - m]
+            dp[n, m] = b_nm * p[n, m - 1] + c_nm * p[n, m + 1]
+        d_nm = 0.5 * sqrt[2] * sqrt[n]
+        dp[n, n] = d_nm * p[n, n - 1]


Test patch:
diff --git a/harmonica/tests/test_legendre.py b/harmonica/tests/test_legendre.py
new file mode 100644
index 0000000..2968d62
--- /dev/null
+++ b/harmonica/tests/test_legendre.py
@@ -0,0 +1,203 @@
+# Copyright (c) 2018 The Harmonica Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# This code is part of the Fatiando a Terra project (https://www.fatiando.org)
+#
+"""
+Test the associated Legendre function calculations.
+"""
+import math
+
+import numpy as np
+import pytest
+
+from .._spherical_harmonics.legendre import (
+    associated_legendre,
+    associated_legendre_derivative,
+    associated_legendre_full,
+    associated_legendre_full_derivative,
+    associated_legendre_schmidt,
+    associated_legendre_schmidt_derivative,
+)
+
+
+def legendre_analytical(x):
+    "Analytical expressions for unnormalized Legendre functions"
+    max_degree = 4
+    p = np.full((max_degree + 1, max_degree + 1), np.nan)
+    p[0, 0] = 1
+    p[1, 0] = x
+    p[1, 1] = np.sqrt(1 - x**2)
+    p[2, 0] = 1 / 2 * (3 * x**2 - 1)
+    p[2, 1] = 3 * x * np.sqrt(1 - x**2)
+    p[2, 2] = 3 * (1 - x**2)
+    p[3, 0] = 1 / 2 * (5 * x**3 - 3 * x)
+    p[3, 1] = -3 / 2 * (1 - 5 * x**2) * np.sqrt(1 - x**2)
+    p[3, 2] = 15 * x * (1 - x**2)
+    p[3, 3] = 15 * np.sqrt(1 - x**2) ** 3
+    p[4, 0] = 1 / 8 * (35 * x**4 - 30 * x**2 + 3)
+    p[4, 1] = 5 / 2 * (7 * x**3 - 3 * x) * np.sqrt(1 - x**2)
+    p[4, 2] = 15 / 2 * (7 * x**2 - 1) * (1 - x**2)
+    p[4, 3] = 105 * x * np.sqrt(1 - x**2) ** 3
+    p[4, 4] = 105 * np.sqrt(1 - x**2) ** 4
+    return p
+
+
+def legendre_derivative_analytical(x):
+    """
+    Analytical expressions for theta derivatives of unnormalized Legendre
+    functions
+    """
+    max_degree = 4
+    dp = np.full((max_degree + 1, max_degree + 1), np.nan)
+    cos = x
+    sin = np.sqrt(1 - x**2)
+    dp[0, 0] = 0
+    dp[1, 0] = -sin
+    dp[1, 1] = cos
+    dp[2, 0] = -3 * cos * sin
+    dp[2, 1] = 3 * (cos**2 - sin**2)
+    dp[2, 2] = 6 * sin * cos
+    dp[3, 0] = (3 * sin - 15 * cos**2 * sin) / 2
+    dp[3, 1] = 3 / 2 * (5 * cos**3 - 10 * cos * sin**2 - cos)
+    dp[3, 2] = 30 * cos**2 * sin - 15 * sin**3
+    dp[3, 3] = 45 * sin**2 * cos
+    dp[4, 0] = (15 * cos - 35 * cos**3) * sin / 2
+    dp[4, 1] = (
+        35 * cos**4 - 15 * cos**2 + 15 * sin**2 - 105 * cos**2 * sin**2
+    ) / 2
+    dp[4, 2] = 105 * sin * cos**3 - 105 * cos * sin**3 - 15 * sin * cos
+    dp[4, 3] = 315 * cos**2 * sin**2 - 105 * sin**4
+    dp[4, 4] = 420 * sin**3 * cos
+    return dp
+
+
+def schmidt_normalization(p):
+    "Multiply by the Schmidt normalization factor"
+    max_degree = p.shape[0] - 1
+    for n in range(max_degree + 1):
+        for m in range(n + 1):
+            if m > 0:
+                p[n, m] *= np.sqrt(2 * math.factorial(n - m) / math.factorial(n + m))
+
+
+def full_normalization(p):
+    "Multiply by the full normalization factor"
+    max_degree = p.shape[0] - 1
+    for n in range(max_degree + 1):
+        for m in range(n + 1):
+            if m == 0:
+                p[n, m] *= np.sqrt(2 * n + 1)
+            else:
+                p[n, m] *= np.sqrt(
+                    2 * (2 * n + 1) * math.factorial(n - m) / math.factorial(n + m)
+                )
+
+
+@pytest.mark.parametrize(
+    "func,norm",
+    (
+        (associated_legendre, None),
+        (associated_legendre_schmidt, schmidt_normalization),
+        (associated_legendre_full, full_normalization),
+    ),
+    ids=["unnormalized", "schmidt", "full"],
+)
+def test_associated_legendre_function_analytical(func, norm):
+    "Check if the first few degrees match analytical expressions"
+    for angle in np.linspace(0, np.pi, 360):
+        x = np.cos(angle)
+        # Analytical expression
+        p_analytical = legendre_analytical(x)
+        max_degree = p_analytical.shape[0] - 1
+        if norm is not None:
+            norm(p_analytical)
+        # Numerical calculation
+        p = np.empty((max_degree + 1, max_degree + 1))
+        func(x, max_degree, p)
+        for n in range(max_degree + 1):
+            for m in range(n + 1):
+                np.testing.assert_allclose(p_analytical[n, m], p[n, m], atol=1e-10)
+
+
+@pytest.mark.parametrize(
+    "func,deriv,norm",
+    (
+        (associated_legendre, associated_legendre_derivative, None),
+        (
+            associated_legendre_schmidt,
+            associated_legendre_schmidt_derivative,
+            schmidt_normalization,
+        ),
+        (
+            associated_legendre_full,
+            associated_legendre_full_derivative,
+            full_normalization,
+        ),
+    ),
+    ids=["unnormalized", "schmidt", "full"],
+)
+def test_associated_legendre_function_derivative_analytical(func, norm, deriv):
+    "Check if the first few degrees match analytical expressions"
+    for angle in np.linspace(0, np.pi, 360):
+        x = np.cos(angle)
+        # Analytical expression
+        dp_analytical = legendre_derivative_analytical(x)
+        max_degree = dp_analytical.shape[0] - 1
+        if norm is not None:
+            norm(dp_analytical)
+        # Numerical calculation
+        p = np.empty((max_degree + 1, max_degree + 1))
+        dp = np.empty((max_degree + 1, max_degree + 1))
+        func(x, max_degree, p)
+        deriv(max_degree, p, dp)
+        for n in range(max_degree + 1):
+            for m in range(n + 1):
+                np.testing.assert_allclose(
+                    dp_analytical[n, m], dp[n, m], atol=1e-10, err_msg=f"n={n}, m={m}"
+                )
+
+
+def test_associated_legengre_function_schmidt_identity():
+    "Check Schmidt normalized functions against a known identity"
+    # Higher degrees than this yield bad results
+    max_degree = 2800
+    # The sum of the coefs squared for a degree should be 1
+    true_value = np.ones(max_degree + 1)
+    p = np.zeros((max_degree + 1, max_degree + 1))
+    for x in np.linspace(-1, 1, 50):
+        associated_legendre_schmidt(x, max_degree, p)
+        np.testing.assert_allclose((p**2).sum(axis=1), true_value, atol=1e-10, rtol=0)
+
+
+# Not testing unnormalized ones because they only works until a very low degree
+@pytest.mark.parametrize(
+    "func,deriv",
+    (
+        (associated_legendre_schmidt, associated_legendre_schmidt_derivative),
+        (associated_legendre_full, associated_legendre_full_derivative),
+    ),
+    ids=["schmidt", "full"],
+)
+def test_associated_legengre_function_legendre_equation(func, deriv):
+    "Check functions and derivatives against the Legendre equation"
+    max_degree = 2800
+    # Legendre equation should result in 0
+    true_value = np.zeros((max_degree + 1, max_degree + 1))
+    p = np.zeros((max_degree + 1, max_degree + 1))
+    dp = np.zeros((max_degree + 1, max_degree + 1))
+    dp2 = np.zeros((max_degree + 1, max_degree + 1))
+    index = np.arange(max_degree + 1).reshape((max_degree + 1, 1))
+    n = np.repeat(index, max_degree + 1, axis=1)
+    m = np.repeat(index.T, max_degree + 1, axis=0)
+    for angle in np.linspace(0.001, np.pi - 0.001, 10):
+        cos = np.cos(angle)
+        sin = np.sin(angle)
+        func(cos, max_degree, p)
+        deriv(max_degree, p, dp)
+        deriv(max_degree, dp, dp2)
+        legendre = sin * dp2 + cos * dp + (sin * n * (n + 1) - m**2 / sin) * p
+        np.testing.assert_allclose(
+            legendre, true_value, atol=1e-5, rtol=0, err_msg=f"angle={angle}"
+        )