quadrature.py 7.8 KB
Newer Older
Boris Bonev's avatar
Boris Bonev committed
1
2
3
4
# coding=utf-8

# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
5
#
Boris Bonev's avatar
Boris Bonev committed
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
# 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 Tuple, Optional
from torch_harmonics.cache import lru_cache
import math
Boris Bonev's avatar
Boris Bonev committed
35
import numpy as np
Thorsten Kurth's avatar
Thorsten Kurth committed
36
import torch
Boris Bonev's avatar
Boris Bonev committed
37

Thorsten Kurth's avatar
Thorsten Kurth committed
38
39
def _precompute_grid(n: int, grid: Optional[str]="equidistant", a: Optional[float]=0.0, b: Optional[float]=1.0,
                     periodic: Optional[bool]=False) -> Tuple[torch.Tensor, torch.Tensor]:
Boris Bonev's avatar
Boris Bonev committed
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

    if (grid != "equidistant") and periodic:
        raise ValueError(f"Periodic grid is only supported on equidistant grids.")

    # compute coordinates
    if grid == "equidistant":
        xlg, wlg = trapezoidal_weights(n, a=a, b=b, periodic=periodic)
    elif grid == "legendre-gauss":
        xlg, wlg = legendre_gauss_weights(n, a=a, b=b)
    elif grid == "lobatto":
        xlg, wlg = lobatto_weights(n, a=a, b=b)
    elif grid == "equiangular":
        xlg, wlg = clenshaw_curtiss_weights(n, a=a, b=b)
    else:
        raise ValueError(f"Unknown grid type {grid}")

    return xlg, wlg

Thorsten Kurth's avatar
Thorsten Kurth committed
58
59
60
61
62
63
64
65
66
@lru_cache(typed=True, copy=True)
def _precompute_longitudes(nlon: int):
    r"""
    Convenience routine to precompute longitudes
    """
    
    lons = torch.linspace(0, 2 * math.pi, nlon+1, dtype=torch.float64, requires_grad=False)[:-1]
    return lons

67

Thorsten Kurth's avatar
Thorsten Kurth committed
68
69
@lru_cache(typed=True, copy=True)
def _precompute_latitudes(nlat: int, grid: Optional[str]="equiangular") -> Tuple[torch.Tensor, torch.Tensor]:
70
71
72
    r"""
    Convenience routine to precompute latitudes
    """
Thorsten Kurth's avatar
Thorsten Kurth committed
73
        
74
    # compute coordinates in the cosine theta domain
Boris Bonev's avatar
Boris Bonev committed
75
    xlg, wlg = _precompute_grid(nlat, grid=grid, a=-1.0, b=1.0, periodic=False)
Thorsten Kurth's avatar
Thorsten Kurth committed
76
    
77
78
    # to perform the quadrature and account for the jacobian of the sphere, the quadrature rule
    # is formulated in the cosine theta domain, which is designed to integrate functions of cos theta
Thorsten Kurth's avatar
Thorsten Kurth committed
79
80
81
    lats = torch.flip(torch.arccos(xlg), dims=(0,)).clone()
    wlg = torch.flip(wlg, dims=(0,)).clone()
    
82
83
    return lats, wlg

84

Thorsten Kurth's avatar
Thorsten Kurth committed
85
def trapezoidal_weights(n: int, a: Optional[float]=-1.0, b: Optional[float]=1.0, periodic: Optional[bool]=False) -> Tuple[torch.Tensor, torch.Tensor]:
Boris Bonev's avatar
Boris Bonev committed
86
87
88
89
90
    r"""
    Helper routine which returns equidistant nodes with trapezoidal weights
    on the interval [a, b]
    """

Thorsten Kurth's avatar
Thorsten Kurth committed
91
    xlg = torch.as_tensor(np.linspace(a, b, n, endpoint=periodic))
Thorsten Kurth's avatar
Thorsten Kurth committed
92
    wlg = (b - a) / (n - periodic * 1) * torch.ones(n, requires_grad=False)
Boris Bonev's avatar
Boris Bonev committed
93
94
95
96
97
98
99

    if not periodic:
        wlg[0] *= 0.5
        wlg[-1] *= 0.5

    return xlg, wlg

100

Thorsten Kurth's avatar
Thorsten Kurth committed
101
def legendre_gauss_weights(n: int, a: Optional[float]=-1.0, b: Optional[float]=1.0) -> Tuple[torch.Tensor, torch.Tensor]:
102
    r"""
Boris Bonev's avatar
Boris Bonev committed
103
104
105
106
107
    Helper routine which returns the Legendre-Gauss nodes and weights
    on the interval [a, b]
    """

    xlg, wlg = np.polynomial.legendre.leggauss(n)
Thorsten Kurth's avatar
Thorsten Kurth committed
108
109
    xlg = torch.as_tensor(xlg).clone()
    wlg = torch.as_tensor(wlg).clone()
Boris Bonev's avatar
Boris Bonev committed
110
111
112
113
114
    xlg = (b - a) * 0.5 * xlg + (b + a) * 0.5
    wlg = wlg * (b - a) * 0.5

    return xlg, wlg

115

Thorsten Kurth's avatar
Thorsten Kurth committed
116
117
def lobatto_weights(n: int, a: Optional[float]=-1.0, b: Optional[float]=1.0,
                    tol: Optional[float]=1e-16, maxiter: Optional[int]=100) -> Tuple[torch.Tensor, torch.Tensor]:
118
    r"""
Boris Bonev's avatar
Boris Bonev committed
119
120
121
122
    Helper routine which returns the Legendre-Gauss-Lobatto nodes and weights
    on the interval [a, b]
    """

Thorsten Kurth's avatar
Thorsten Kurth committed
123
124
125
    wlg = torch.zeros((n,), dtype=torch.float64, requires_grad=False)
    tlg = torch.zeros((n,), dtype=torch.float64, requires_grad=False)
    tmp = torch.zeros((n,), dtype=torch.float64, requires_grad=False)
Boris Bonev's avatar
Boris Bonev committed
126
127

    # Vandermonde Matrix
Thorsten Kurth's avatar
Thorsten Kurth committed
128
    vdm = torch.zeros((n, n), dtype=torch.float64, requires_grad=False)
129

Boris Bonev's avatar
Boris Bonev committed
130
    # initialize Chebyshev nodes as first guess
131
    for i in range(n):
132
        tlg[i] = -math.cos(math.pi * i / (n - 1))
133

Boris Bonev's avatar
Boris Bonev committed
134
    tmp = 2.0
135

Boris Bonev's avatar
Boris Bonev committed
136
137
    for i in range(maxiter):
        tmp = tlg
138
139
140
141

        vdm[:, 0] = 1.0
        vdm[:, 1] = tlg

Boris Bonev's avatar
Boris Bonev committed
142
        for k in range(2, n):
143
144
145
146
147
148
149
150
            vdm[:, k] = ((2 * k - 1) * tlg * vdm[:, k - 1] - (k - 1) * vdm[:, k - 2]) / k

        tlg = tmp - (tlg * vdm[:, n - 1] - vdm[:, n - 2]) / (n * vdm[:, n - 1])

        if max(abs(tlg - tmp).flatten()) < tol:
            break

    wlg = 2.0 / ((n * (n - 1)) * (vdm[:, n - 1] ** 2))
Boris Bonev's avatar
Boris Bonev committed
151
152
153
154

    # rescale
    tlg = (b - a) * 0.5 * tlg + (b + a) * 0.5
    wlg = wlg * (b - a) * 0.5
155

Boris Bonev's avatar
Boris Bonev committed
156
157
158
    return tlg, wlg


Thorsten Kurth's avatar
Thorsten Kurth committed
159
def clenshaw_curtiss_weights(n: int, a: Optional[float]=-1.0, b: Optional[float]=1.0) -> Tuple[torch.Tensor, torch.Tensor]:
160
    r"""
Boris Bonev's avatar
Boris Bonev committed
161
162
163
164
165
166
    Computation of the Clenshaw-Curtis quadrature nodes and weights.
    This implementation follows

    [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018.
    """

167
    assert n > 1
Boris Bonev's avatar
Boris Bonev committed
168

Thorsten Kurth's avatar
Thorsten Kurth committed
169
    tcc = torch.cos(torch.linspace(math.pi, 0, n, dtype=torch.float64, requires_grad=False))
Boris Bonev's avatar
Boris Bonev committed
170
171

    if n == 2:
Thorsten Kurth's avatar
Thorsten Kurth committed
172
        wcc = torch.tensor([1.0, 1.0], dtype=torch.float64)
Boris Bonev's avatar
Boris Bonev committed
173
174
175
    else:

        n1 = n - 1
Thorsten Kurth's avatar
Thorsten Kurth committed
176
        N = torch.arange(1, n1, 2, dtype=torch.float64)
Boris Bonev's avatar
Boris Bonev committed
177
178
179
        l = len(N)
        m = n1 - l

Thorsten Kurth's avatar
Thorsten Kurth committed
180
181
182
        v = torch.cat([2 / N / (N - 2), 1 / N[-1:], torch.zeros(m, dtype=torch.float64, requires_grad=False)])
        #v = 0 - v[:-1] - v[-1:0:-1]
        v = 0 - v[:-1] - torch.flip(v[1:], dims=(0,))
Boris Bonev's avatar
Boris Bonev committed
183

Thorsten Kurth's avatar
Thorsten Kurth committed
184
        g0 = -torch.ones(n1, dtype=torch.float64, requires_grad=False)
Boris Bonev's avatar
Boris Bonev committed
185
186
        g0[l] = g0[l] + n1
        g0[m] = g0[m] + n1
187
        g = g0 / (n1**2 - 1 + (n1 % 2))
Thorsten Kurth's avatar
Thorsten Kurth committed
188
189
        wcc = torch.fft.ifft(v + g).real
        wcc = torch.cat((wcc, wcc[:1]))
Boris Bonev's avatar
Boris Bonev committed
190
191
192
193
194
195
196

    # rescale
    tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5
    wcc = wcc * (b - a) * 0.5

    return tcc, wcc

197

Thorsten Kurth's avatar
Thorsten Kurth committed
198
def fejer2_weights(n: int, a: Optional[float]=-1.0, b: Optional[float]=1.0) -> Tuple[torch.Tensor, torch.Tensor]:
199
    r"""
Boris Bonev's avatar
Boris Bonev committed
200
201
202
203
204
205
    Computation of the Fejer quadrature nodes and weights.
    This implementation follows

    [1] Joerg Waldvogel, Fast Construction of the Fejer and Clenshaw-Curtis Quadrature Rules; BIT Numerical Mathematics, Vol. 43, No. 1, pp. 001–018.
    """

206
    assert n > 2
Boris Bonev's avatar
Boris Bonev committed
207

Thorsten Kurth's avatar
Thorsten Kurth committed
208
    tcc = torch.cos(torch.linspace(math.pi, 0, n, dtype=torch.float64, requires_grad=False))
Boris Bonev's avatar
Boris Bonev committed
209
210

    n1 = n - 1
Thorsten Kurth's avatar
Thorsten Kurth committed
211
    N = torch.arange(1, n1, 2, dtype=torch.float64)
Boris Bonev's avatar
Boris Bonev committed
212
213
214
    l = len(N)
    m = n1 - l

Thorsten Kurth's avatar
Thorsten Kurth committed
215
216
217
    v = torch.cat([2 / N / (N - 2), 1 / N[-1:], torch.zeros(m, dtype=torch.float64, requires_grad=False)])
    #v = 0 - v[:-1] - v[-1:0:-1]
    v = 0 - v[:-1] - torch.flip(v[1:], dims=(0,))
Boris Bonev's avatar
Boris Bonev committed
218

Thorsten Kurth's avatar
Thorsten Kurth committed
219
220
    wcc = torch.fft.ifft(v).real
    wcc = torch.cat((wcc, wcc[:1]))
Boris Bonev's avatar
Boris Bonev committed
221
222
223
224
225

    # rescale
    tcc = (b - a) * 0.5 * tcc + (b + a) * 0.5
    wcc = wcc * (b - a) * 0.5

226
    return tcc, wcc