test_convolution.py 13.7 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
32
33
34
35
36
37
38
39
40
# 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.
#

import unittest
from parameterized import parameterized
from functools import partial
import math
import numpy as np
import torch
from torch.autograd import gradcheck
from torch_harmonics import *

Boris Bonev's avatar
Boris Bonev committed
41
from torch_harmonics.quadrature import _precompute_grid, _precompute_latitudes
Boris Bonev's avatar
Boris Bonev committed
42

43

Boris Bonev's avatar
Boris Bonev committed
44
def _compute_vals_isotropic(r: torch.Tensor, phi: torch.Tensor, nr: int, r_cutoff: float):
Boris Bonev's avatar
Boris Bonev committed
45
46
47
48
    """
    helper routine to compute the values of the isotropic kernel densely
    """

Boris Bonev's avatar
Boris Bonev committed
49
50
    kernel_size = (nr // 2) + nr % 2
    ikernel = torch.arange(kernel_size).reshape(-1, 1, 1)
51
    dr = 2 * r_cutoff / (nr + 1)
Boris Bonev's avatar
Boris Bonev committed
52

Boris Bonev's avatar
Boris Bonev committed
53
    # compute the support
Boris Bonev's avatar
Boris Bonev committed
54
55
56
57
    if nr % 2 == 1:
        ir = ikernel * dr
    else:
        ir = (ikernel + 0.5) * dr
Boris Bonev's avatar
Boris Bonev committed
58
59

    vals = torch.where(
Boris Bonev's avatar
Boris Bonev committed
60
61
        ((r - ir).abs() <= dr) & (r <= r_cutoff),
        (1 - (r - ir).abs() / dr),
Boris Bonev's avatar
Boris Bonev committed
62
63
        0,
    )
Boris Bonev's avatar
Boris Bonev committed
64

Boris Bonev's avatar
Boris Bonev committed
65
66
    return vals

Boris Bonev's avatar
Boris Bonev committed
67
68

def _compute_vals_anisotropic(r: torch.Tensor, phi: torch.Tensor, nr: int, nphi: int, r_cutoff: float):
Boris Bonev's avatar
Boris Bonev committed
69
70
71
72
    """
    helper routine to compute the values of the anisotropic kernel densely
    """

Boris Bonev's avatar
Boris Bonev committed
73
    kernel_size = (nr // 2) * nphi + nr % 2
Boris Bonev's avatar
Boris Bonev committed
74
    ikernel = torch.arange(kernel_size).reshape(-1, 1, 1)
75
    dr = 2 * r_cutoff / (nr + 1)
Boris Bonev's avatar
Boris Bonev committed
76
    dphi = 2.0 * math.pi / nphi
Boris Bonev's avatar
Boris Bonev committed
77

Boris Bonev's avatar
Boris Bonev committed
78
79
80
81
82
83
84
85
86
87
88
89
90
    # disambiguate even and uneven cases and compute the support
    if nr % 2 == 1:
        ir = ((ikernel - 1) // nphi + 1) * dr
        iphi = ((ikernel - 1) % nphi) * dphi
    else:
        ir = (ikernel // nphi + 0.5) * dr
        iphi = (ikernel % nphi) * dphi

    # compute the value of the filter
    if nr % 2 == 1:
        # find the indices where the rotated position falls into the support of the kernel
        cond_r = ((r - ir).abs() <= dr) & (r <= r_cutoff)
        cond_phi = ((phi - iphi).abs() <= dphi) | ((2 * math.pi - (phi - iphi).abs()) <= dphi)
91
        r_vals = torch.where(cond_r, (1 - (r - ir).abs() / dr), 0.0)
Boris Bonev's avatar
Boris Bonev committed
92
93
94
95
96
97
98
99
100
101
102
        phi_vals = torch.where(cond_phi, (1 - torch.minimum((phi - iphi).abs(), (2 * math.pi - (phi - iphi).abs())) / dphi), 0.0)
        vals = torch.where(ikernel > 0, r_vals * phi_vals, r_vals)
    else:
        # find the indices where the rotated position falls into the support of the kernel
        cond_r = ((r - ir).abs() <= dr) & (r <= r_cutoff)
        cond_phi = ((phi - iphi).abs() <= dphi) | ((2 * math.pi - (phi - iphi).abs()) <= dphi)
        r_vals = torch.where(cond_r, (1 - (r - ir).abs() / dr), 0.0)
        phi_vals = torch.where(cond_phi, (1 - torch.minimum((phi - iphi).abs(), (2 * math.pi - (phi - iphi).abs())) / dphi), 0.0)
        vals = r_vals * phi_vals

        # in the even case, the inner casis functions overlap into areas with a negative areas
103
104
        rn = -r
        phin = torch.where(phi + math.pi >= 2 * math.pi, phi - math.pi, phi + math.pi)
Boris Bonev's avatar
Boris Bonev committed
105
106
107
108
109
        cond_rn = ((rn - ir).abs() <= dr) & (rn <= r_cutoff)
        cond_phin = ((phin - iphi).abs() <= dphi) | ((2 * math.pi - (phin - iphi).abs()) <= dphi)
        rn_vals = torch.where(cond_rn, (1 - (rn - ir).abs() / dr), 0.0)
        phin_vals = torch.where(cond_phin, (1 - torch.minimum((phin - iphi).abs(), (2 * math.pi - (phin - iphi).abs())) / dphi), 0.0)
        vals += rn_vals * phin_vals
Boris Bonev's avatar
Boris Bonev committed
110
111
112

    return vals

113

114
def _normalize_convolution_tensor_dense(psi, quad_weights, transpose_normalization=False, merge_quadrature=False, eps=1e-9):
Boris Bonev's avatar
Boris Bonev committed
115
116
117
118
119
120
121
122
123
124
    """
    Discretely normalizes the convolution tensor.
    """

    kernel_size, nlat_out, nlon_out, nlat_in, nlon_in = psi.shape
    scale_factor = float(nlon_in // nlon_out)

    if transpose_normalization:
        # the normalization is not quite symmetric due to the compressed way psi is stored in the main code
        # look at the normalization code in the actual implementation
125
        psi_norm = torch.sum(quad_weights.reshape(1, -1, 1, 1, 1) * psi[:, :, :1], dim=(1, 4), keepdim=True) / scale_factor
126
127
        if merge_quadrature:
            psi = quad_weights.reshape(1, -1, 1, 1, 1) * psi
Boris Bonev's avatar
Boris Bonev committed
128
129
    else:
        psi_norm = torch.sum(quad_weights.reshape(1, 1, 1, -1, 1) * psi, dim=(3, 4), keepdim=True)
130
131
        if merge_quadrature:
            psi = quad_weights.reshape(1, 1, 1, -1, 1) * psi
Boris Bonev's avatar
Boris Bonev committed
132
133
134
135

    return psi / (psi_norm + eps)


136
137
138
139
140
141
142
143
144
145
146
def _precompute_convolution_tensor_dense(
    in_shape,
    out_shape,
    kernel_shape,
    quad_weights,
    grid_in="equiangular",
    grid_out="equiangular",
    theta_cutoff=0.01 * math.pi,
    transpose_normalization=False,
    merge_quadrature=False,
):
Boris Bonev's avatar
Boris Bonev committed
147
148
149
150
151
152
153
    """
    Helper routine to compute the convolution Tensor in a dense fashion
    """

    assert len(in_shape) == 2
    assert len(out_shape) == 2

Boris Bonev's avatar
Boris Bonev committed
154
155
    quad_weights = quad_weights.reshape(-1, 1)

Boris Bonev's avatar
Boris Bonev committed
156
    if len(kernel_shape) == 1:
Boris Bonev's avatar
Boris Bonev committed
157
        kernel_handle = partial(_compute_vals_isotropic, nr=kernel_shape[0], r_cutoff=theta_cutoff)
158
        kernel_size = math.ceil(kernel_shape[0] / 2)
Boris Bonev's avatar
Boris Bonev committed
159
    elif len(kernel_shape) == 2:
Boris Bonev's avatar
Boris Bonev committed
160
161
        kernel_handle = partial(_compute_vals_anisotropic, nr=kernel_shape[0], nphi=kernel_shape[1], r_cutoff=theta_cutoff)
        kernel_size = (kernel_shape[0] // 2) * kernel_shape[1] + kernel_shape[0] % 2
Boris Bonev's avatar
Boris Bonev committed
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
    else:
        raise ValueError("kernel_shape should be either one- or two-dimensional.")

    nlat_in, nlon_in = in_shape
    nlat_out, nlon_out = out_shape

    lats_in, _ = quadrature._precompute_latitudes(nlat_in, grid=grid_in)
    lats_in = torch.from_numpy(lats_in).float()
    lats_out, _ = quadrature._precompute_latitudes(nlat_out, grid=grid_out)
    lats_out = torch.from_numpy(lats_out).float()  # array for accumulating non-zero indices

    # compute the phi differences. We need to make the linspace exclusive to not double the last point
    lons_in = torch.linspace(0, 2 * math.pi, nlon_in + 1)[:-1]
    lons_out = torch.linspace(0, 2 * math.pi, nlon_out + 1)[:-1]

    out = torch.zeros(kernel_size, nlat_out, nlon_out, nlat_in, nlon_in)

    for t in range(nlat_out):
        for p in range(nlon_out):
            alpha = -lats_out[t]
            beta = lons_in - lons_out[p]
            gamma = lats_in.reshape(-1, 1)

            # compute latitude of the rotated position
            z = -torch.cos(beta) * torch.sin(alpha) * torch.sin(gamma) + torch.cos(alpha) * torch.cos(gamma)

            # compute cartesian coordinates of the rotated position
            x = torch.cos(alpha) * torch.cos(beta) * torch.sin(gamma) + torch.cos(gamma) * torch.sin(alpha)
            y = torch.sin(beta) * torch.sin(gamma)

            # normalize instead of clipping to ensure correct range
            norm = torch.sqrt(x * x + y * y + z * z)
            x = x / norm
            y = y / norm
            z = z / norm

            # compute spherical coordinates
            theta = torch.arccos(z)
            phi = torch.arctan2(y, x) + torch.pi

            # find the indices where the rotated position falls into the support of the kernel
            out[:, t, p, :, :] = kernel_handle(theta, phi)

Boris Bonev's avatar
Boris Bonev committed
205
    # take care of normalization
206
    out = _normalize_convolution_tensor_dense(out, quad_weights=quad_weights, transpose_normalization=transpose_normalization, merge_quadrature=merge_quadrature)
Boris Bonev's avatar
Boris Bonev committed
207

Boris Bonev's avatar
Boris Bonev committed
208
209
210
211
212
213
    return out


class TestDiscreteContinuousConvolution(unittest.TestCase):
    def setUp(self):
        if torch.cuda.is_available():
214
215
216
            self.device = torch.device("cuda:0")
            torch.cuda.set_device(self.device.index)
            torch.cuda.manual_seed(333)
Boris Bonev's avatar
Boris Bonev committed
217
218
        else:
            self.device = torch.device("cpu")
Boris Bonev's avatar
Boris Bonev committed
219

220
        torch.manual_seed(333)
Boris Bonev's avatar
Boris Bonev committed
221

Boris Bonev's avatar
Boris Bonev committed
222
223
224
    @parameterized.expand(
        [
            # regular convolution
Boris Bonev's avatar
Boris Bonev committed
225
226
227
228
229
230
231
232
233
            [8, 4, 2, (16, 32), (16, 32), [3], "equiangular", "equiangular", False, 1e-4],
            [8, 4, 2, (16, 32), (8, 16), [5], "equiangular", "equiangular", False, 1e-4],
            [8, 4, 2, (16, 32), (8, 16), [3, 3], "equiangular", "equiangular", False, 1e-4],
            [8, 4, 2, (16, 32), (8, 16), [4, 3], "equiangular", "equiangular", False, 1e-4],
            [8, 4, 2, (16, 24), (8, 8), [3], "equiangular", "equiangular", False, 1e-4],
            [8, 4, 2, (18, 36), (6, 12), [7], "equiangular", "equiangular", False, 1e-4],
            [8, 4, 2, (16, 32), (8, 16), [5], "equiangular", "legendre-gauss", False, 1e-4],
            [8, 4, 2, (16, 32), (8, 16), [5], "legendre-gauss", "equiangular", False, 1e-4],
            [8, 4, 2, (16, 32), (8, 16), [5], "legendre-gauss", "legendre-gauss", False, 1e-4],
Boris Bonev's avatar
Boris Bonev committed
234
            # transpose convolution
Boris Bonev's avatar
Boris Bonev committed
235
236
237
238
239
240
241
242
243
            [8, 4, 2, (16, 32), (16, 32), [3], "equiangular", "equiangular", True, 1e-4],
            [8, 4, 2, (8, 16), (16, 32), [5], "equiangular", "equiangular", True, 1e-4],
            [8, 4, 2, (8, 16), (16, 32), [3, 3], "equiangular", "equiangular", True, 1e-4],
            [8, 4, 2, (8, 16), (16, 32), [4, 3], "equiangular", "equiangular", True, 1e-4],
            [8, 4, 2, (8, 8), (16, 24), [3], "equiangular", "equiangular", True, 1e-4],
            [8, 4, 2, (6, 12), (18, 36), [7], "equiangular", "equiangular", True, 1e-4],
            [8, 4, 2, (8, 16), (16, 32), [5], "equiangular", "legendre-gauss", True, 1e-4],
            [8, 4, 2, (8, 16), (16, 32), [5], "legendre-gauss", "equiangular", True, 1e-4],
            [8, 4, 2, (8, 16), (16, 32), [5], "legendre-gauss", "legendre-gauss", True, 1e-4],
Boris Bonev's avatar
Boris Bonev committed
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
        ]
    )
    def test_disco_convolution(
        self,
        batch_size,
        in_channels,
        out_channels,
        in_shape,
        out_shape,
        kernel_shape,
        grid_in,
        grid_out,
        transpose,
        tol,
    ):
Boris Bonev's avatar
Boris Bonev committed
259
260
261
262
263
        nlat_in, nlon_in = in_shape
        nlat_out, nlon_out = out_shape

        theta_cutoff = (kernel_shape[0] + 1) / 2 * torch.pi / float(nlat_out - 1)

Boris Bonev's avatar
Boris Bonev committed
264
        Conv = DiscreteContinuousConvTransposeS2 if transpose else DiscreteContinuousConvS2
265
266
267
        conv = Conv(in_channels, out_channels, in_shape, out_shape, kernel_shape, groups=1, grid_in=grid_in, grid_out=grid_out, bias=False, theta_cutoff=theta_cutoff).to(
            self.device
        )
Boris Bonev's avatar
Boris Bonev committed
268

Boris Bonev's avatar
Boris Bonev committed
269
270
        _, wgl = _precompute_latitudes(nlat_in, grid=grid_in)
        quad_weights = 2.0 * torch.pi * torch.from_numpy(wgl).float().reshape(-1, 1) / nlon_in
Boris Bonev's avatar
Boris Bonev committed
271
272

        if transpose:
273
274
275
            psi_dense = _precompute_convolution_tensor_dense(
                out_shape, in_shape, kernel_shape, quad_weights, grid_in=grid_out, grid_out=grid_in, theta_cutoff=theta_cutoff, transpose_normalization=True, merge_quadrature=True
            ).to(self.device)
Boris Bonev's avatar
Boris Bonev committed
276
277
278
279

            psi = torch.sparse_coo_tensor(conv.psi_idx, conv.psi_vals, size=(conv.kernel_size, conv.nlat_in, conv.nlat_out * conv.nlon_out)).to_dense()

            self.assertTrue(torch.allclose(psi, psi_dense[:, :, 0].reshape(-1, nlat_in, nlat_out * nlon_out)))
Boris Bonev's avatar
Boris Bonev committed
280
        else:
281
282
283
            psi_dense = _precompute_convolution_tensor_dense(
                in_shape, out_shape, kernel_shape, quad_weights, grid_in=grid_in, grid_out=grid_out, theta_cutoff=theta_cutoff, transpose_normalization=False, merge_quadrature=True
            ).to(self.device)
Boris Bonev's avatar
Boris Bonev committed
284

Boris Bonev's avatar
Boris Bonev committed
285
286
            psi = torch.sparse_coo_tensor(conv.psi_idx, conv.psi_vals, size=(conv.kernel_size, conv.nlat_out, conv.nlat_in * conv.nlon_in)).to_dense()

Boris Bonev's avatar
Boris Bonev committed
287
            self.assertTrue(torch.allclose(psi, psi_dense[:, :, 0].reshape(-1, nlat_out, nlat_in * nlon_in)))
Boris Bonev's avatar
Boris Bonev committed
288
289

        # create a copy of the weight
290
291
292
293
        w_ref = torch.empty_like(conv.weight)
        with torch.no_grad():
            w_ref.copy_(conv.weight)
        w_ref.requires_grad = True
Boris Bonev's avatar
Boris Bonev committed
294
295

        # create an input signal
296
297
298
299
300
301
302
303
        x = torch.randn(batch_size, in_channels, *in_shape, device=self.device)

        # FWD and BWD pass
        x.requires_grad = True
        y = conv(x)
        grad_input = torch.randn_like(y)
        y.backward(grad_input)
        x_grad = x.grad.clone()
Boris Bonev's avatar
Boris Bonev committed
304
305
306

        # perform the reference computation
        x_ref = x.clone().detach()
307
        x_ref.requires_grad = True
Boris Bonev's avatar
Boris Bonev committed
308
309
        if transpose:
            y_ref = torch.einsum("oif,biqr->bofqr", w_ref, x_ref)
310
            y_ref = torch.einsum("fqrtp,bofqr->botp", psi_dense, y_ref)
Boris Bonev's avatar
Boris Bonev committed
311
        else:
312
            y_ref = torch.einsum("ftpqr,bcqr->bcftp", psi_dense, x_ref)
Boris Bonev's avatar
Boris Bonev committed
313
            y_ref = torch.einsum("oif,biftp->botp", w_ref, y_ref)
314
        y_ref.backward(grad_input)
Boris Bonev's avatar
Boris Bonev committed
315
316
        x_ref_grad = x_ref.grad.clone()

Boris Bonev's avatar
Boris Bonev committed
317
318
319
        # compare results
        self.assertTrue(torch.allclose(y, y_ref, rtol=tol, atol=tol))

Boris Bonev's avatar
Boris Bonev committed
320
        # compare
321
        self.assertTrue(torch.allclose(x_grad, x_ref_grad, rtol=tol, atol=tol))
Boris Bonev's avatar
Boris Bonev committed
322
323
        self.assertTrue(torch.allclose(conv.weight.grad, w_ref.grad, rtol=tol, atol=tol))

Boris Bonev's avatar
Boris Bonev committed
324

Boris Bonev's avatar
Boris Bonev committed
325
326
if __name__ == "__main__":
    unittest.main()