diff --git a/harmonica/_spherical_harmonics/legendre.py b/harmonica/_spherical_harmonics/legendre.py new file mode 100644 index 000000000..7e25a3d61 --- /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] diff --git a/harmonica/tests/test_legendre.py b/harmonica/tests/test_legendre.py new file mode 100644 index 000000000..2968d6246 --- /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}" + )