small fixes after rebase

Author: azrael417

tests/testutils.py

@@ -0,0 +1,44 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 torch
+
+def compare_tensors(msg, tensor, tensor_ref, rtol=1e-8, atol=1e-5,  verbose=False):
+    allclose = torch.allclose(tensor, tensor_ref, rtol=rtol, atol=atol)
+    if (not allclose) and verbose:
+        diff = torch.abs(tensor - tensor_ref)
+        print(f"{msg} absolute tensor diff: min = {torch.min(diff)}, mean = {torch.mean(diff)}, max = {torch.max(diff)}.")
+        reldiff = diff / torch.abs(tensor_ref)
+        print(f"{msg} relative tensor diff: min = {torch.min(reldiff)}, mean = {torch.mean(reldiff)}, max = {torch.max(reldiff)}.")
+        # find element with maximum difference
+        index = torch.argmax(diff)
+        print(f"{msg} element {index} with maximum difference: value = {tensor.flatten()[index]}, reference value = {tensor_ref.flatten()[index]}, diff = {diff.flatten()[index]}.")
+    return allclose

torch_harmonics/disco/convolution.py

@@ -370,8 +370,10 @@ def __init__(
             raise ValueError("Error, the number of output channels has to be an integer multiple of the group size")
         self.groupsize_in = in_channels // self.groups
         self.groupsize_out = out_channels // self.groups
+        # keep this for backward compatibility
+        self.groupsize = self.groupsize_in
         scale = math.sqrt(1.0 / self.groupsize_in / self.kernel_size)
-        self.weight = nn.Parameter(scale * torch.randn(self.groups, self.groupsize_out, self.groupsize_in, self.kernel_size))
+        self.weight = nn.Parameter(scale * torch.randn(self.groups * self.groupsize_out, self.groupsize_in, self.kernel_size))
 
         if bias:
             self.bias = nn.Parameter(torch.zeros(out_channels))
@@ -496,7 +498,7 @@ def __init__(
             self.psi = _get_psi(self.kernel_size, self.psi_idx, self.psi_vals, self.nlat_in, self.nlon_in, self.nlat_out, self.nlon_out)
 
     def extra_repr(self):
-        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize_in * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
 
     @property
     def psi_idx(self):
@@ -524,7 +526,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             # weight multiplication
             B, H, W, _, K = xpc.shape
             xpc = xpc.reshape(B, H, W, self.groups, self.groupsize_in, K)
-            outp = torch.einsum("bxygck,gock->bxygo", xpc, self.weight)
+            outp = torch.einsum("bxygck,gock->bxygo", xpc, self.weight.reshape(self.groups, self.groupsize_out, self.groupsize_in, self.kernel_size))
             outp = outp.reshape(B, H, W, -1).contiguous()
 
             # permute output
@@ -538,7 +540,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = x.reshape(B, self.groups, self.groupsize_in, K, H, W)
 
             # weight multiplication
-            out = torch.einsum("bgckxy,gock->bgoxy", x, self.weight)
+            out = torch.einsum("bgckxy,gock->bgoxy", x, self.weight.reshape(self.groups, self.groupsize_out, self.groupsize_in, self.kernel_size))
             out = out.reshape(B, -1, H, W).contiguous()
 
         if self.bias is not None:
@@ -657,7 +659,7 @@ def __init__(
             self.psi_st = _get_psi(self.kernel_size, self.psi_idx, self.psi_vals, self.nlat_in, self.nlon_in, self.nlat_out, self.nlon_out, semi_transposed=True)
 
     def extra_repr(self):
-        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize_in * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
 
     @property
     def psi_idx(self):
@@ -674,7 +676,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
             # weight multiplication
             xp = xp.reshape(B, H, W, self.groups, self.groupsize_in)
-            xpc = torch.einsum("bxygc,gock->bxygok", xp, self.weight)
+            xpc = torch.einsum("bxygc,gock->bxygok", xp, self.weight.reshape(self.groups, self.groupsize_out, self.groupsize_in, self.kernel_size))
             xpc = xpc.reshape(B, H, W, -1, self.kernel_size).contiguous()
 
             # disco contraction
@@ -695,7 +697,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         else:
             # weight multiplication
             x = x.reshape(B, self.groups, self.groupsize_in, H, W)
-            xc = torch.einsum("bgcxy,gock->bgokxy", x, self.weight)
+            xc = torch.einsum("bgcxy,gock->bgokxy", x, self.weight.reshape(self.groups, self.groupsize_out, self.groupsize_in, self.kernel_size))
             xc = xc.reshape(B, self.groups* self.groupsize_out, -1, H, W).contiguous()
 
             # disco contraction