legendre.py 9.01 KB
Newer Older
Boris Bonev's avatar
Boris Bonev committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
# coding=utf-8

# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

Thorsten Kurth's avatar
Thorsten Kurth committed
32
33
34
from typing import Optional
import math
import torch
Boris Bonev's avatar
Boris Bonev committed
35

Thorsten Kurth's avatar
Thorsten Kurth committed
36
37
38
39
from torch_harmonics.cache import lru_cache


def clm(l: int, m: int) -> float:
Boris Bonev's avatar
Boris Bonev committed
40
41
    """
    defines the normalization factor to orthonormalize the Spherical Harmonics
apaaris's avatar
apaaris committed
42
43
44
45
46
47
48
49
50
51
52
53

    Parameters
    -----------
    l: int
        Degree of the spherical harmonic
    m: int
        Order of the spherical harmonic

    Returns
    -------
    out: float
        Normalization factor
Boris Bonev's avatar
Boris Bonev committed
54
    """
Thorsten Kurth's avatar
Thorsten Kurth committed
55
    return math.sqrt((2*l + 1) / 4 / math.pi) * math.sqrt(math.factorial(l-m) / math.factorial(l+m))
Boris Bonev's avatar
Boris Bonev committed
56

Thorsten Kurth's avatar
Thorsten Kurth committed
57
def legpoly(mmax: int, lmax: int, x: torch.Tensor, norm: Optional[str]="ortho", inverse: Optional[bool]=False, csphase: Optional[bool]=True) -> torch.Tensor:
58
    r"""
59
60
61
    Computes the values of (-1)^m c^l_m P^l_m(x) at the positions specified by x.
    The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m
    can be turned off optionally.
Boris Bonev's avatar
Boris Bonev committed
62

apaaris's avatar
apaaris committed
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    Parameters
    -----------
    mmax: int
        Maximum order of the spherical harmonics
    lmax: int
        Maximum degree of the spherical harmonics
    x: torch.Tensor
        Tensor of positions at which to evaluate the Legendre polynomials
    norm: Optional[str]
        Normalization of the Legendre polynomials
    inverse: Optional[bool]
        Whether to compute the inverse Legendre polynomials
    csphase: Optional[bool]
        Whether to apply the Condon-Shortley phase (-1)^m

    Returns
    -------
    out: torch.Tensor
        Tensor of Legendre polynomial values

    References
    ----------
Boris Bonev's avatar
Boris Bonev committed
85
86
87
88
89
90
91
    [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
    [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982;
        https://apps.dtic.mil/sti/citations/ADA123406
    [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients
    """

    # compute the tensor P^m_n:
Boris Bonev's avatar
Boris Bonev committed
92
    nmax = max(mmax,lmax)
Thorsten Kurth's avatar
Thorsten Kurth committed
93
    vdm = torch.zeros((nmax, nmax, len(x)), dtype=torch.float64, device=x.device, requires_grad=False)
Boris Bonev's avatar
Boris Bonev committed
94
        
Thorsten Kurth's avatar
Thorsten Kurth committed
95
96
    norm_factor = 1.0 if norm == "ortho" else math.sqrt(4 * math.pi)
    norm_factor = 1.0 / norm_factor if inverse else norm_factor
Boris Bonev's avatar
Boris Bonev committed
97
98

    # initial values to start the recursion
Thorsten Kurth's avatar
Thorsten Kurth committed
99
    vdm[0,0,:] = norm_factor / math.sqrt(4 * math.pi)
Boris Bonev's avatar
Boris Bonev committed
100
101

    # fill the diagonal and the lower diagonal
Boris Bonev's avatar
Boris Bonev committed
102
    for l in range(1, nmax):
Thorsten Kurth's avatar
Thorsten Kurth committed
103
104
        vdm[l-1, l, :] = math.sqrt(2*l + 1) * x * vdm[l-1, l-1, :]
        vdm[l, l, :] = torch.sqrt( (2*l + 1) * (1 + x) * (1 - x) / 2 / l ) * vdm[l-1, l-1, :]
Boris Bonev's avatar
Boris Bonev committed
105
106

    # fill the remaining values on the upper triangle and multiply b
Boris Bonev's avatar
Boris Bonev committed
107
    for l in range(2, nmax):
Boris Bonev's avatar
Boris Bonev committed
108
        for m in range(0, l-1):
Thorsten Kurth's avatar
Thorsten Kurth committed
109
110
            vdm[m, l, :] = x * math.sqrt((2*l - 1) / (l - m) * (2*l + 1) / (l + m)) * vdm[m, l-1, :] \
                            - math.sqrt((l + m - 1) / (l - m) * (2*l + 1) / (2*l - 3) * (l - m - 1) / (l + m)) * vdm[m, l-2, :]
Boris Bonev's avatar
Boris Bonev committed
111
112

    if norm == "schmidt":
Boris Bonev's avatar
Boris Bonev committed
113
        for l in range(0, nmax):
Boris Bonev's avatar
Boris Bonev committed
114
            if inverse:
Thorsten Kurth's avatar
Thorsten Kurth committed
115
                vdm[:, l, : ] = vdm[:, l, : ] * math.sqrt(2*l + 1)
Boris Bonev's avatar
Boris Bonev committed
116
            else:
Thorsten Kurth's avatar
Thorsten Kurth committed
117
                vdm[:, l, : ] = vdm[:, l, : ] / math.sqrt(2*l + 1)
Boris Bonev's avatar
Boris Bonev committed
118

119
    vdm = vdm[:mmax, :lmax]
Boris Bonev's avatar
Boris Bonev committed
120

Boris Bonev's avatar
Boris Bonev committed
121
122
    if csphase:
        for m in range(1, mmax, 2):
123
124
125
126
            vdm[m] *= -1

    return vdm

Thorsten Kurth's avatar
Thorsten Kurth committed
127
128
129
@lru_cache(typed=True, copy=True)
def _precompute_legpoly(mmax: int , lmax: int, t: torch.Tensor,
                        norm: Optional[str]="ortho", inverse: Optional[bool]=False, csphase: Optional[bool]=True) -> torch.Tensor:
130
131
132
133
134
    r"""
    Computes the values of (-1)^m c^l_m P^l_m(\cos \theta) at the positions specified by t (theta).
    The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m
    can be turned off optionally.

apaaris's avatar
apaaris committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
    Parameters
    -----------
    mmax: int
        Maximum order of the spherical harmonics
    lmax: int
        Maximum degree of the spherical harmonics
    t: torch.Tensor
        Tensor of positions at which to evaluate the Legendre polynomials
    norm: Optional[str]
        Normalization of the Legendre polynomials
    inverse: Optional[bool]
        Whether to compute the inverse Legendre polynomials
    csphase: Optional[bool]
        Whether to apply the Condon-Shortley phase (-1)^m

    Returns
    -------
    out: torch.Tensor
        Tensor of Legendre polynomial values

    References
    ----------
157
158
159
160
161
    [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
    [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982;
        https://apps.dtic.mil/sti/citations/ADA123406
    [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients
    """
Boris Bonev's avatar
Boris Bonev committed
162

Thorsten Kurth's avatar
Thorsten Kurth committed
163
    return legpoly(mmax, lmax, torch.cos(t), norm=norm, inverse=inverse, csphase=csphase)
Boris Bonev's avatar
Boris Bonev committed
164

Thorsten Kurth's avatar
Thorsten Kurth committed
165
166
167
@lru_cache(typed=True, copy=True)
def _precompute_dlegpoly(mmax: int, lmax: int, t: torch.Tensor,
                         norm: Optional[str]="ortho", inverse: Optional[bool]=False, csphase: Optional[bool]=True) -> torch.Tensor:
168
    r"""
Boris Bonev's avatar
Boris Bonev committed
169
    Computes the values of the derivatives $\frac{d}{d \theta} P^m_l(\cos \theta)$
170
    at the positions specified by t (theta), as well as $\frac{1}{\sin \theta} P^m_l(\cos \theta)$,
Boris Bonev's avatar
Boris Bonev committed
171
    needed for the computation of the vector spherical harmonics. The resulting tensor has shape
172
    (2, mmax, lmax, len(t)).
Boris Bonev's avatar
Boris Bonev committed
173

apaaris's avatar
apaaris committed
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
    Parameters
    -----------
    mmax: int
        Maximum order of the spherical harmonics
    lmax: int
        Maximum degree of the spherical harmonics
    t: torch.Tensor
        Tensor of positions at which to evaluate the Legendre polynomials
    norm: Optional[str]
        Normalization of the Legendre polynomials
    inverse: Optional[bool]
        Whether to compute the inverse Legendre polynomials
    csphase: Optional[bool]
        Whether to apply the Condon-Shortley phase (-1)^m

    Returns
    -------
    out: torch.Tensor
        Tensor of Legendre polynomial values

    References
    ----------
Boris Bonev's avatar
Boris Bonev committed
196
197
198
    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
    """

199
    pct = _precompute_legpoly(mmax+1, lmax+1, t, norm=norm, inverse=inverse, csphase=False)
Boris Bonev's avatar
Boris Bonev committed
200

201
    dpct = torch.zeros((2, mmax, lmax, len(t)), dtype=torch.float64, device=t.device, requires_grad=False)
Boris Bonev's avatar
Boris Bonev committed
202
203
204
205
206

    # fill the derivative terms wrt theta
    for l in range(0, lmax):

        # m = 0
Thorsten Kurth's avatar
Thorsten Kurth committed
207
        dpct[0, 0, l] = - math.sqrt(l*(l+1)) * pct[1, l]
Boris Bonev's avatar
Boris Bonev committed
208
209
210

        # 0 < m < l
        for m in range(1, min(l, mmax)):
Thorsten Kurth's avatar
Thorsten Kurth committed
211
            dpct[0, m, l] = 0.5 * ( math.sqrt((l+m)*(l-m+1)) * pct[m-1, l] - math.sqrt((l-m)*(l+m+1)) * pct[m+1, l] )
Boris Bonev's avatar
Boris Bonev committed
212
213
214

        # m == l
        if mmax > l:
Thorsten Kurth's avatar
Thorsten Kurth committed
215
            dpct[0, l, l] = math.sqrt(l/2) * pct[l-1, l]
Boris Bonev's avatar
Boris Bonev committed
216
217
218
219
220
221

        # fill the - 1j m P^m_l / sin(phi). as this component is purely imaginary,
        # we won't store it explicitly in a complex array
        for m in range(1, min(l+1, mmax)):
            # this component is implicitly complex
            # we do not divide by m here as this cancels with the derivative of the exponential
Thorsten Kurth's avatar
Thorsten Kurth committed
222
223
            dpct[1, m, l] = 0.5 * math.sqrt((2*l+1)/(2*l+3)) * \
                ( math.sqrt((l-m+1)*(l-m+2)) * pct[m-1, l+1] + math.sqrt((l+m+1)*(l+m+2)) * pct[m+1, l+1] )
Boris Bonev's avatar
Boris Bonev committed
224
225
226
227
228

    if csphase:
        for m in range(1, mmax, 2):
            dpct[:, m] *= -1

Thorsten Kurth's avatar
Thorsten Kurth committed
229
    return dpct