test_convolution.py 13.5 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

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

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

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

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

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

Boris Bonev's avatar
Boris Bonev committed
66
67

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

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

Boris Bonev's avatar
Boris Bonev committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
    # 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)
        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 = 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
        rn = - r
        phin = torch.where(phi + math.pi >= 2*math.pi, phi - math.pi, phi + math.pi)
        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
109
110
111

    return vals

Boris Bonev's avatar
Boris Bonev committed
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def _normalize_convolution_tensor_dense(psi, quad_weights, transpose_normalization=False, eps=1e-9):
    """
    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
        psi_norm = torch.sum(quad_weights.reshape(1, -1, 1, 1, 1) * psi[:,:,:1], dim=(1, 4), keepdim=True) / scale_factor
    else:
        psi_norm = torch.sum(quad_weights.reshape(1, 1, 1, -1, 1) * psi, dim=(3, 4), keepdim=True)

    return psi / (psi_norm + eps)


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):
Boris Bonev's avatar
Boris Bonev committed
131
132
133
134
135
136
137
    """
    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
138
139
    quad_weights = quad_weights.reshape(-1, 1)

Boris Bonev's avatar
Boris Bonev committed
140
    if len(kernel_shape) == 1:
Boris Bonev's avatar
Boris Bonev committed
141
142
        kernel_handle = partial(_compute_vals_isotropic, nr=kernel_shape[0], r_cutoff=theta_cutoff)
        kernel_size = math.ceil( kernel_shape[0] / 2)
Boris Bonev's avatar
Boris Bonev committed
143
    elif len(kernel_shape) == 2:
Boris Bonev's avatar
Boris Bonev committed
144
145
        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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
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
    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
189
190
191
    # take care of normalization
    out = _normalize_convolution_tensor_dense(out, quad_weights=quad_weights, transpose_normalization=transpose_normalization)

Boris Bonev's avatar
Boris Bonev committed
192
193
194
195
196
197
    return out


class TestDiscreteContinuousConvolution(unittest.TestCase):
    def setUp(self):
        if torch.cuda.is_available():
198
199
200
            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
201
202
        else:
            self.device = torch.device("cpu")
Boris Bonev's avatar
Boris Bonev committed
203

204
        torch.manual_seed(333)
Boris Bonev's avatar
Boris Bonev committed
205

Boris Bonev's avatar
Boris Bonev committed
206
207
208
    @parameterized.expand(
        [
            # regular convolution
Boris Bonev's avatar
Boris Bonev committed
209
210
211
212
213
214
215
216
217
            [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
218
            # transpose convolution
Boris Bonev's avatar
Boris Bonev committed
219
220
221
222
223
224
225
226
227
            [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
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
        ]
    )
    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
243
244
245
246
247
        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
248
249
250
251
252
253
254
255
256
257
258
        Conv = DiscreteContinuousConvTransposeS2 if transpose else DiscreteContinuousConvS2
        conv = Conv(
            in_channels,
            out_channels,
            in_shape,
            out_shape,
            kernel_shape,
            groups=1,
            grid_in=grid_in,
            grid_out=grid_out,
            bias=False,
Boris Bonev's avatar
Boris Bonev committed
259
            theta_cutoff=theta_cutoff
Boris Bonev's avatar
Boris Bonev committed
260
261
        ).to(self.device)

Boris Bonev's avatar
Boris Bonev committed
262
263
        _, 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
264
265

        if transpose:
Boris Bonev's avatar
Boris Bonev committed
266
267
268
269
270
            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).to(self.device)

            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
271
        else:
Boris Bonev's avatar
Boris Bonev committed
272
            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).to(self.device)
Boris Bonev's avatar
Boris Bonev committed
273

Boris Bonev's avatar
Boris Bonev committed
274
275
            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
276
            self.assertTrue(torch.allclose(psi, psi_dense[:, :, 0].reshape(-1, nlat_out, nlat_in * nlon_in)))
Boris Bonev's avatar
Boris Bonev committed
277
278

        # create a copy of the weight
279
280
281
282
        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
283
284

        # create an input signal
285
286
287
288
289
290
291
292
        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
293
294
295

        # perform the reference computation
        x_ref = x.clone().detach()
296
        x_ref.requires_grad = True
Boris Bonev's avatar
Boris Bonev committed
297
298
299
300
301
302
        if transpose:
            y_ref = torch.einsum("oif,biqr->bofqr", w_ref, x_ref)
            y_ref = torch.einsum("fqrtp,bofqr->botp", psi_dense, y_ref * conv.quad_weights)
        else:
            y_ref = torch.einsum("ftpqr,bcqr->bcftp", psi_dense, x_ref * conv.quad_weights)
            y_ref = torch.einsum("oif,biftp->botp", w_ref, y_ref)
303
        y_ref.backward(grad_input)
Boris Bonev's avatar
Boris Bonev committed
304
305
        x_ref_grad = x_ref.grad.clone()

Boris Bonev's avatar
Boris Bonev committed
306
307
308
        # compare results
        self.assertTrue(torch.allclose(y, y_ref, rtol=tol, atol=tol))

Boris Bonev's avatar
Boris Bonev committed
309
        # compare
310
        self.assertTrue(torch.allclose(x_grad, x_ref_grad, rtol=tol, atol=tol))
Boris Bonev's avatar
Boris Bonev committed
311
312
        self.assertTrue(torch.allclose(conv.weight.grad, w_ref.grad, rtol=tol, atol=tol))

Boris Bonev's avatar
Boris Bonev committed
313

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