# 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 abc
import math
import warnings
from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
from disco_helpers import optimized_kernels_is_available, pack_psi_dense, preprocess_psi
from torch_harmonics.cache import lru_cache
from torch_harmonics.filter_basis import FilterBasis, get_filter_basis
from torch_harmonics.quadrature import precompute_latitudes, precompute_longitudes
from ._disco_utils import _get_psi
from .kernels_torch.disco_torch import _disco_s2_contraction_torch, _disco_s2_transpose_contraction_torch
from .optimized.disco_optimized import (
_build_kernel_split_csr,
_disco_s2_contraction_kpacked,
_disco_s2_contraction_optimized,
_disco_s2_conv_save_x_kpacked,
_disco_s2_conv_save_x_optimized,
_disco_s2_fused_conv_kpacked,
_disco_s2_fused_conv_optimized,
_disco_s2_transpose_contraction_optimized,
_kpacked_supported_on_device,
_maybe_kpack_psi,
_split_csr_python_offsets,
_use_spatial_first_dgrad,
)
def _kpacked_device_supported_for_tensor(tensor: torch.Tensor) -> bool:
if not tensor.is_cuda:
return False
return _kpacked_supported_on_device(tensor.get_device())
def _normalize_convolution_tensor_s2(
psi_idx,
psi_vals,
in_shape,
out_shape,
kernel_size,
quad_weights,
theta_cutoff,
transpose_normalization=False,
basis_norm_mode="mean",
merge_quadrature=False,
isotropic_mask=None,
eps=1e-9,
):
"""Normalizes convolution tensor values based on specified normalization mode.
This function applies different normalization strategies to the convolution tensor
values based on the basis_norm_mode parameter. It can normalize individual basis
functions, compute mean normalization across all basis functions, or use support
weights. The function also optionally merges quadrature weights into the tensor.
The implementation is fully vectorized: each nonzero is assigned a flat group id
``gid = ikernel * nlat_out + ilat_out`` and all per-group sums (support, bias
numerator, scale) are accumulated in a single ``scatter_add_`` per reduced
quantity, avoiding the per-(ikernel, ilat_out) Python double loop.
Parameters
----------
psi_idx : torch.Tensor
Index tensor for the sparse convolution tensor.
psi_vals : torch.Tensor
Value tensor for the sparse convolution tensor.
in_shape : Tuple[int]
Tuple of (nlat_in, nlon_in) representing input grid dimensions.
out_shape : Tuple[int]
Tuple of (nlat_out, nlon_out) representing output grid dimensions.
kernel_size : int
Number of kernel basis functions.
quad_weights : torch.Tensor
Quadrature weights for numerical integration.
theta_cutoff : float
Angular cutoff of the filter support (radians). Required by the "geometric" mode,
which normalizes by the theoretical area measure of the spherical cap of half-angle
theta_cutoff; unused by other modes.
transpose_normalization : bool
If True, applies normalization in transpose direction.
basis_norm_mode : str
Normalization mode, one of ["none", "nodal", "modal", "mean", "support", "geometric"].
The legacy names "individual" and "area ratio" are accepted as deprecated aliases
for "nodal" and "geometric" respectively; each emits a DeprecationWarning.
merge_quadrature : bool
If True, multiplies values by quadrature weights.
isotropic_mask : Optional[Sequence[bool]]
Per-kernel-index boolean mask; True marks an axisymmetric (m=0) basis function.
Used by the "modal" mode to decide which kernels get a weighted-mean bias
subtraction (anisotropic only). If None, only kernel index 0 is treated as isotropic.
eps : float
Small epsilon value to prevent division by zero.
Returns
-------
torch.Tensor
Normalized convolution tensor values.
Raises
------
ValueError
If basis_norm_mode is not one of the supported modes.
"""
if basis_norm_mode == "individual":
warnings.warn(
'basis_norm_mode="individual" is deprecated, use "nodal" instead.',
DeprecationWarning,
stacklevel=2,
)
basis_norm_mode = "nodal"
elif basis_norm_mode == "area ratio":
warnings.warn(
'basis_norm_mode="area ratio" is deprecated, use "geometric" instead.',
DeprecationWarning,
stacklevel=2,
)
basis_norm_mode = "geometric"
# reshape the indices implicitly to be ikernel, out_shape[0], in_shape[0], in_shape[1]
idx = torch.stack([psi_idx[0], psi_idx[1], psi_idx[2] // in_shape[1], psi_idx[2] % in_shape[1]], dim=0)
ikernel = idx[0]
if transpose_normalization:
ilat_out = idx[2]
ilat_in = idx[1]
# deliberately swap input/output shapes to handle transpose normalization with the same code
nlat_out = in_shape[0]
correction_factor = out_shape[1] / in_shape[1]
else:
ilat_out = idx[1]
ilat_in = idx[2]
nlat_out = out_shape[0]
# quadrature weight per nonzero
q = quad_weights[ilat_in].reshape(-1)
# group id per nonzero: (ikernel, ilat_out) -> flat index in [0, kernel_size * nlat_out)
n_groups = kernel_size * nlat_out
gid = ikernel * nlat_out + ilat_out
# support[ik, ilat] = sum_{nonzeros in group} q
support_flat = torch.zeros(n_groups, dtype=psi_vals.dtype, device=psi_vals.device)
support_flat.scatter_add_(0, gid, q)
support = support_flat.view(kernel_size, nlat_out)
# bias[ik, ilat] -- only nonzero for "modal" mode on anisotropic kernels.
# Quadrature-weighted mean of psi_vals over the (ik, ilat) neighborhood.
bias = torch.zeros(kernel_size, nlat_out, dtype=psi_vals.dtype, device=psi_vals.device)
if basis_norm_mode == "modal":
if isotropic_mask is not None:
iso = torch.as_tensor(isotropic_mask, dtype=torch.bool, device=psi_vals.device)
else:
iso = torch.zeros(kernel_size, dtype=torch.bool, device=psi_vals.device)
iso[0] = True
aniso_per_nz = (~iso)[ikernel].to(psi_vals.dtype)
bias_num_flat = torch.zeros(n_groups, dtype=psi_vals.dtype, device=psi_vals.device)
bias_num_flat.scatter_add_(0, gid, psi_vals * q * aniso_per_nz)
# divide; isotropic kernels have bias_num=0 so result stays 0; clamp protects empty groups
bias = (bias_num_flat / support_flat.clamp(min=eps)).view(kernel_size, nlat_out)
# zero the bias on empty groups
bias = torch.where(support.abs() > eps, bias, torch.zeros_like(bias))
# scale[ik, ilat] = sum |psi_vals - bias[ik, ilat]| * q over the neighborhood
bias_per_nz = bias.view(-1)[gid]
scale_flat = torch.zeros(n_groups, dtype=psi_vals.dtype, device=psi_vals.device)
scale_flat.scatter_add_(0, gid, (psi_vals - bias_per_nz).abs() * q)
scale = scale_flat.view(kernel_size, nlat_out)
# per-mode (b, s) selection per nonzero, then renormalize in a single elementwise pass
if basis_norm_mode in ("nodal", "modal"):
scale_per_nz = scale.view(-1)[gid]
psi_vals = (psi_vals - bias_per_nz) / scale_per_nz.clamp(min=eps)
elif basis_norm_mode == "mean":
# average over latitudes per kernel; bias is zero in this mode
bias_per_ik = bias.mean(dim=1)
scale_per_ik = scale.mean(dim=1)
bias_per_nz_per_ik = bias_per_ik[ikernel]
scale_per_nz_per_ik = scale_per_ik[ikernel]
psi_vals = (psi_vals - bias_per_nz_per_ik) / scale_per_nz_per_ik.clamp(min=eps)
elif basis_norm_mode == "support":
support_scale_per_nz = support.view(-1)[gid]
psi_vals = psi_vals / support_scale_per_nz.clamp(min=eps)
elif basis_norm_mode == "geometric":
geometric_scale = (1.0 - math.cos(theta_cutoff)) / 2.0 / 2.0
psi_vals = psi_vals / max(geometric_scale, eps)
elif basis_norm_mode == "none":
pass
else:
raise ValueError(f"Unknown basis normalization mode {basis_norm_mode}.")
if merge_quadrature:
psi_vals = psi_vals * q
if transpose_normalization and merge_quadrature:
psi_vals = psi_vals / correction_factor
return psi_vals
@lru_cache(typed=True, copy=True)
def _precompute_convolution_tensor_s2(
in_shape: Tuple[int],
out_shape: Tuple[int],
filter_basis: FilterBasis,
grid_in: Optional[str] = "equiangular",
grid_out: Optional[str] = "equiangular",
theta_cutoff: Optional[float] = 0.01 * math.pi,
theta_eps: Optional[float] = 1e-3,
transpose_normalization: Optional[bool] = False,
basis_norm_mode: Optional[str] = "nodal",
merge_quadrature: Optional[bool] = False,
):
r"""
Precomputes the rotated filters at positions $R^{-1}_j \omega_i = R^{-1}_j R_i \nu = Y(-\theta_j)Z(\phi_i - \phi_j)Y(\theta_j)\nu$.
Assumes a tensorized grid on the sphere with an equiangular-trapezoidal sampling in longitude as described in Ocampo et al.
The output tensor has shape kernel_shape x nlat_out x (nlat_in * nlon_in).
The rotation of the Euler angles uses the YZY convention, which applied to the northpole $(0,0,1)^T$ yields
$$
Y(\alpha) Z(\beta) Y(\gamma) n =
{\begin{bmatrix}
\cos(\gamma)\sin(\alpha) + \cos(\alpha)\cos(\beta)\sin(\gamma) \\
\sin(\beta)\sin(\gamma) \\
\cos(\alpha)\cos(\gamma)-\cos(\beta)\sin(\alpha)\sin(\gamma)
\end{bmatrix}}
$$
Parameters
----------
in_shape : Tuple[int]
Input shape of the convolution tensor
out_shape : Tuple[int]
Output shape of the convolution tensor
filter_basis : FilterBasis
Filter basis functions
grid_in : str
Input grid type
grid_out : str
Output grid type
theta_cutoff : float
Theta cutoff for the filter basis functions
theta_eps : float
Epsilon for the theta cutoff
transpose_normalization : bool
Whether to normalize the convolution tensor in the transpose direction
basis_norm_mode : str
Mode for basis normalization
merge_quadrature : bool
Whether to merge the quadrature weights into the convolution tensor
Returns
-------
out_idx : torch.Tensor
Index tensor of the convolution tensor
out_vals : torch.Tensor
Values tensor of the convolution tensor
"""
if len(in_shape) != 2:
raise ValueError(f"in_shape must be a 2-tuple (nlat, nlon), got length {len(in_shape)}")
if len(out_shape) != 2:
raise ValueError(f"out_shape must be a 2-tuple (nlat, nlon), got length {len(out_shape)}")
kernel_size = filter_basis.kernel_size
nlat_in, nlon_in = in_shape
nlat_out, nlon_out = out_shape
# precompute input and output grids
lats_in, win = precompute_latitudes(nlat_in, grid=grid_in)
lats_out, wout = precompute_latitudes(nlat_out, grid=grid_out)
# compute the phi differences
# It's imporatant to not include the 2 pi point in the longitudes, as it is equivalent to lon=0
lons_in = precompute_longitudes(nlon_in)
# compute quadrature weights and merge them into the convolution tensor.
# These quadrature integrate to 1 over the sphere.
if transpose_normalization:
quad_weights = wout.reshape(-1, 1) / nlon_in / 2.0
else:
quad_weights = win.reshape(-1, 1) / nlon_in / 2.0
# effective theta cutoff if multiplied with a fudge factor to avoid aliasing with grid width (especially near poles)
theta_cutoff_eff = (1.0 + theta_eps) * theta_cutoff
out_idx = []
out_vals = []
beta = lons_in
gamma = lats_in.reshape(-1, 1)
# compute trigs
cbeta = torch.cos(beta)
sbeta = torch.sin(beta)
cgamma = torch.cos(gamma)
sgamma = torch.sin(gamma)
# compute row offsets
out_roff = torch.zeros(nlat_out + 1, dtype=torch.int64, device=lons_in.device)
out_roff[0] = 0
for t in range(nlat_out):
# the last angle has a negative sign as it is a passive rotation, which rotates the filter around the y-axis
alpha = -lats_out[t]
# compute cartesian coordinates of the rotated position
# This uses the YZY convention of Euler angles, where the last angle (alpha) is a passive rotation,
# and therefore applied with a negative sign
x = torch.cos(alpha) * cbeta * sgamma + cgamma * torch.sin(alpha)
y = sbeta * sgamma
z = -cbeta * torch.sin(alpha) * sgamma + torch.cos(alpha) * cgamma
# normalization is important to avoid NaNs when arccos and atan are applied
# this can otherwise lead to spurious artifacts in the solution
norm = torch.sqrt(x * x + y * y + z * z)
x = x / norm
y = y / norm
z = z / norm
# compute spherical coordinates, where phi needs to fall into the [0, 2pi) range
theta = torch.arccos(z)
phi = torch.arctan2(y, x)
phi = torch.where(phi < 0.0, phi + 2 * torch.pi, phi)
# find the indices where the rotated position falls into the support of the kernel
iidx, vals = filter_basis.compute_support_vals(theta, phi, r_cutoff=theta_cutoff_eff)
# add the output latitude and reshape such that psi has dimensions kernel_shape x nlat_out x (nlat_in*nlon_in)
idx = torch.stack([iidx[:, 0], t * torch.ones_like(iidx[:, 0]), iidx[:, 1] * nlon_in + iidx[:, 2]], dim=0)
# append indices and values to the COO datastructure, compute row offsets
out_idx.append(idx)
out_vals.append(vals)
out_roff[t + 1] = out_roff[t] + iidx.shape[0]
# concatenate the indices and values
out_idx = torch.cat(out_idx, dim=-1)
out_vals = torch.cat(out_vals, dim=-1)
out_vals = _normalize_convolution_tensor_s2(
out_idx,
out_vals,
in_shape,
out_shape,
kernel_size,
quad_weights,
theta_cutoff,
transpose_normalization=transpose_normalization,
basis_norm_mode=basis_norm_mode,
merge_quadrature=merge_quadrature,
isotropic_mask=filter_basis.isotropic_mask,
)
out_idx = out_idx.contiguous()
out_vals = out_vals.contiguous()
return out_idx, out_vals, out_roff
class DiscreteContinuousConv(nn.Module, metaclass=abc.ABCMeta):
"""
Abstract base class for discrete-continuous convolutions
Parameters
----------
in_channels : int
Number of input channels
out_channels : int
Number of output channels
kernel_shape : Union[int, Tuple[int], Tuple[int, int]]
Shape of the kernel
basis_type : Optional[str]
Type of the basis functions
groups : Optional[int]
Number of groups
bias : Optional[bool]
Whether to use bias
Returns
-------
torch.Tensor
Output tensor
"""
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
basis_type: Optional[str] = "piecewise linear",
groups: Optional[int] = 1,
bias: Optional[bool] = True,
optimized_kernel: Optional[bool] = True,
):
super().__init__()
self.kernel_shape = kernel_shape
self.optimized_kernel = optimized_kernel and optimized_kernels_is_available()
# get the filter basis functions
self.filter_basis = get_filter_basis(kernel_shape=kernel_shape, basis_type=basis_type)
# groups
self.groups = groups
# weight tensor
if in_channels % self.groups != 0:
raise ValueError("Error, the number of input channels has to be an integer multiple of the group size")
if out_channels % self.groups != 0:
raise ValueError("Error, the number of output channels has to be an integer multiple of the group size")
self.groupsize = in_channels // self.groups
self.out_per_group = out_channels // self.groups
scale = math.sqrt(1.0 / self.groupsize) * self.filter_basis.get_init_factors().reshape(1, 1, -1)
self.weight = nn.Parameter(scale * torch.randn(out_channels, self.groupsize, self.kernel_size))
if bias:
self.bias = nn.Parameter(torch.zeros(out_channels))
else:
self.bias = None
@property
def kernel_size(self):
return self.filter_basis.kernel_size
@abc.abstractmethod
def forward(self, x: torch.Tensor):
raise NotImplementedError
[docs]
class DiscreteContinuousConvS2(DiscreteContinuousConv):
r"""
Discrete-continuous (DISCO) convolution on the 2-sphere, as described in :cite:`Ocampo2023`.
The layer evaluates a spherical convolution with a compactly supported
filter of angular radius ``theta_cutoff``. The filter is parameterised as
a learnable linear combination of fixed basis functions
:math:`\{\phi_k\}`, and the integral is computed by sparse quadrature over
the input grid, giving :math:`O(N)` cost in the number of grid points.
The forward pass is
.. math::
g^{c_o}(\theta'_j, \lambda'_q)
= \sum_{c_i} \sum_k w_k^{c_o,c_i}
\sum_{i,\,p} \Psi_{k,\,j,\,(i,p)}\;
f^{c_i}(\theta_i, \lambda'_q + \lambda_p)
where :math:`\Psi` is a precomputed sparse convolution tensor that
encodes the basis function values at rotated input grid positions,
weighted by the quadrature weights. Because the grid is equispaced in
longitude, :math:`\Psi` is independent of the output longitude
(p-shift symmetry).
.. seealso::
:doc:`/guide/disco_convolutions`
User guide with the full mathematical derivation, filter basis
visualisations, and worked examples.
Parameters
----------
in_channels : int
Number of input channels
out_channels : int
Number of output channels
in_shape : Tuple[int]
Input shape of the convolution tensor
out_shape : Tuple[int]
Output shape of the convolution tensor
kernel_shape : Union[int, Tuple[int], Tuple[int, int]]
Shape of the kernel
basis_type : Optional[str]
Type of the basis functions
basis_norm_mode : Optional[str]
Mode for basis normalization
groups : Optional[int]
Number of groups
grid_in : Optional[str]
Input grid type
grid_out : Optional[str]
Output grid type
bias : Optional[bool]
Whether to use bias
theta_cutoff : Optional[float]
Theta cutoff for the filter basis functions
optimized_kernel : Optional[bool]
Whether to use the optimized kernel (if available)
fused : Optional[bool]
When True, fuses the sparse contraction and weight multiplication into a single
autograd region to avoid storing the K-expanded intermediate in the graph.
Trades one extra contraction recompute in backward for K× memory savings.
Only effective when optimized_kernel is True.
References
----------
:cite:`Ocampo2023`
"""
def __init__(
self,
in_channels: int,
out_channels: int,
in_shape: Tuple[int],
out_shape: Tuple[int],
kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
basis_type: Optional[str] = "piecewise linear",
basis_norm_mode: Optional[str] = "nodal",
groups: Optional[int] = 1,
grid_in: Optional[str] = "equiangular",
grid_out: Optional[str] = "equiangular",
bias: Optional[bool] = True,
theta_cutoff: Optional[float] = None,
optimized_kernel: Optional[bool] = True,
fused: Optional[bool] = False,
):
super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias, optimized_kernel)
self.fused = fused and self.optimized_kernel
self.nlat_in, self.nlon_in = in_shape
self.nlat_out, self.nlon_out = out_shape
self.kpacked_device_supported = False
# make sure the p-shift works by checking that longitudes are divisible
if self.nlon_in % self.nlon_out != 0:
raise ValueError(f"nlon_in ({self.nlon_in}) must be an integer multiple of nlon_out ({self.nlon_out}) for the DISCO p-shift to be exact")
# heuristic to compute theta cutoff based on the bandlimit of the input field and overlaps of the basis functions
if theta_cutoff is None:
self.theta_cutoff = torch.pi / float(self.nlat_out - 1)
else:
self.theta_cutoff = theta_cutoff
if self.theta_cutoff <= 0.0:
raise ValueError("Error, theta_cutoff has to be positive.")
idx, vals, _ = _precompute_convolution_tensor_s2(
in_shape,
out_shape,
self.filter_basis,
grid_in=grid_in,
grid_out=grid_out,
theta_cutoff=self.theta_cutoff,
transpose_normalization=False,
basis_norm_mode=basis_norm_mode,
merge_quadrature=True,
)
# sort the values
ker_idx = idx[0, ...].contiguous()
row_idx = idx[1, ...].contiguous()
col_idx = idx[2, ...].contiguous()
vals = vals.contiguous()
self.psi_kpacked_K_pad = None # set to int if kpacked buffers are available
if self.optimized_kernel:
# preprocessed data-structure for GPU kernel
roff_idx = preprocess_psi(self.kernel_size, self.nlat_out, ker_idx, row_idx, col_idx, vals).contiguous()
self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
split_roff_idx, split_nnz_off, split_ker_idx, split_row_idx, split_col_idx, split_vals = _build_kernel_split_csr(
roff_idx, ker_idx, row_idx, col_idx, vals, self.kernel_size, self.nlat_out
)
self.psi_split_row_offsets, self.psi_split_nnz_offsets = _split_csr_python_offsets(split_nnz_off)
self.register_buffer("psi_split_roff_idx", split_roff_idx, persistent=False)
self.register_buffer("psi_split_nnz_off", split_nnz_off, persistent=False)
self.register_buffer("psi_split_ker_idx", split_ker_idx, persistent=False)
self.register_buffer("psi_split_row_idx", split_row_idx, persistent=False)
self.register_buffer("psi_split_col_idx", split_col_idx, persistent=False)
self.register_buffer("psi_split_vals", split_vals, persistent=False)
# optional K-packed dense layout for the WGMMA path (Hopper bf16/fp16).
# precompute here so it's available at forward time.
psi_packed_idx, psi_packed_vals, psi_packed_count = pack_psi_dense(
self.kernel_size,
self.nlat_out,
self.nlon_in,
0,
ker_idx,
row_idx,
col_idx,
vals,
roff_idx,
)
kpack = _maybe_kpack_psi(psi_packed_idx.contiguous(), psi_packed_vals.contiguous(), psi_packed_count.contiguous())
if kpack is not None:
kpacked_idx, kpacked_vals, kpacked_count, K_pad = kpack
self.register_buffer("psi_kpacked_idx", kpacked_idx, persistent=False)
self.register_buffer("psi_kpacked_vals", kpacked_vals, persistent=False)
self.register_buffer("psi_kpacked_count", kpacked_count, persistent=False)
self.psi_kpacked_K_pad = K_pad
# optional K-packed dense layout for the WGMMA path (Hopper bf16/fp16).
# precompute here so it's available at forward time.
psi_packed_idx, psi_packed_vals, psi_packed_count = pack_psi_dense(
self.kernel_size,
self.nlat_out,
self.nlon_in,
0,
ker_idx,
row_idx,
col_idx,
vals,
roff_idx,
)
kpack = _maybe_kpack_psi(psi_packed_idx.contiguous(), psi_packed_vals.contiguous(), psi_packed_count.contiguous())
if kpack is not None:
kpacked_idx, kpacked_vals, kpacked_count, K_pad = kpack
self.register_buffer("psi_kpacked_idx", kpacked_idx, persistent=False)
self.register_buffer("psi_kpacked_vals", kpacked_vals, persistent=False)
self.register_buffer("psi_kpacked_count", kpacked_count, persistent=False)
self.psi_kpacked_K_pad = K_pad
# save all datastructures
self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
self.register_buffer("psi_row_idx", row_idx, persistent=False)
self.register_buffer("psi_col_idx", col_idx, persistent=False)
self.register_buffer("psi_vals", vals, persistent=False)
# also store psi as COO matrix just in case for torch input
if not self.optimized_kernel:
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)
# cache static forward-path decisions so the forward sees plain Python bools/ints,
# not symbolic expressions that confuse torch.compile's value-range analysis
self._save_x_spatial_first_ok = self.optimized_kernel and _use_spatial_first_dgrad(self.out_per_group, self.groupsize, self.kernel_size, self.psi_roff_idx, self.nlat_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}, theta_cutoff={self.theta_cutoff}, groups={self.groups}"
@property
def psi_idx(self):
return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
def _refresh_kpacked_device_supported(self):
if not hasattr(self, "psi_vals"):
self.kpacked_device_supported = False
return
self.kpacked_device_supported = _kpacked_device_supported_for_tensor(self.psi_vals)
def _apply(self, fn):
result = super()._apply(fn)
self._refresh_kpacked_device_supported()
return result
[docs]
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Apply the discrete-continuous convolution.
Parameters
----------
x : torch.Tensor
Input signal of shape ``(batch, in_channels, nlat_in, nlon_in)``.
Returns
-------
torch.Tensor
Convolved signal of shape ``(batch, out_channels, nlat_out, nlon_out)``.
"""
weight_r = self.weight.reshape(self.groups, self.out_per_group, self.weight.shape[1], self.weight.shape[2])
kpacked_dtype = x.dtype
if x.is_cuda and torch.is_autocast_enabled("cuda"):
kpacked_dtype = torch.get_autocast_dtype("cuda")
_kpacked_ok = (
self.optimized_kernel
and self.psi_kpacked_K_pad in (8, 16)
and kpacked_dtype in (torch.float16, torch.bfloat16)
and x.is_cuda
and self.nlon_out % 8 == 0
and self.nlon_in % self.nlon_out == 0
and self.kpacked_device_supported
)
_save_x_spatial_first_ok = self._save_x_spatial_first_ok
if self.fused and _kpacked_ok:
out = _disco_s2_fused_conv_kpacked(
x.to(kpacked_dtype),
weight_r,
self.psi_kpacked_idx,
self.psi_kpacked_vals,
self.psi_kpacked_count,
self.psi_roff_idx,
self.psi_ker_idx,
self.psi_row_idx,
self.psi_col_idx,
self.psi_vals,
self.psi_split_roff_idx,
self.psi_split_nnz_off,
self.psi_split_ker_idx,
self.psi_split_row_idx,
self.psi_split_col_idx,
self.psi_split_vals,
self.kernel_size,
self.nlat_out,
self.nlon_out,
self.groups,
self.groupsize,
self.psi_split_row_offsets,
self.psi_split_nnz_offsets,
)
elif self.fused:
out = _disco_s2_fused_conv_optimized(
x,
weight_r,
self.psi_roff_idx,
self.psi_ker_idx,
self.psi_row_idx,
self.psi_col_idx,
self.psi_vals,
self.psi_split_roff_idx,
self.psi_split_nnz_off,
self.psi_split_ker_idx,
self.psi_split_row_idx,
self.psi_split_col_idx,
self.psi_split_vals,
self.kernel_size,
self.nlat_out,
self.nlon_out,
self.groups,
self.groupsize,
self.psi_split_row_offsets,
self.psi_split_nnz_offsets,
)
else:
if _save_x_spatial_first_ok and _kpacked_ok:
out = _disco_s2_conv_save_x_kpacked(
x.to(kpacked_dtype),
weight_r,
self.psi_kpacked_idx,
self.psi_kpacked_vals,
self.psi_kpacked_count,
self.psi_roff_idx,
self.psi_ker_idx,
self.psi_row_idx,
self.psi_col_idx,
self.psi_vals,
self.psi_split_roff_idx,
self.psi_split_nnz_off,
self.psi_split_ker_idx,
self.psi_split_row_idx,
self.psi_split_col_idx,
self.psi_split_vals,
self.kernel_size,
self.nlat_out,
self.nlon_out,
self.groups,
self.groupsize,
self.psi_split_row_offsets,
self.psi_split_nnz_offsets,
)
elif _save_x_spatial_first_ok:
out = _disco_s2_conv_save_x_optimized(
x.to(kpacked_dtype),
weight_r,
self.psi_roff_idx,
self.psi_ker_idx,
self.psi_row_idx,
self.psi_col_idx,
self.psi_vals,
self.psi_split_roff_idx,
self.psi_split_nnz_off,
self.psi_split_ker_idx,
self.psi_split_row_idx,
self.psi_split_col_idx,
self.psi_split_vals,
self.kernel_size,
self.nlat_out,
self.nlon_out,
self.groups,
self.groupsize,
self.psi_split_row_offsets,
self.psi_split_nnz_offsets,
)
elif _kpacked_ok:
x = _disco_s2_contraction_kpacked(
x.to(kpacked_dtype),
self.psi_kpacked_idx,
self.psi_kpacked_vals,
self.psi_kpacked_count,
self.psi_roff_idx,
self.psi_ker_idx,
self.psi_row_idx,
self.psi_col_idx,
self.psi_vals,
self.kernel_size,
self.nlat_out,
self.nlon_out,
)
elif self.optimized_kernel:
x = _disco_s2_contraction_optimized(
x, self.psi_roff_idx, self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx, self.psi_vals, self.kernel_size, self.nlat_out, self.nlon_out
)
else:
x = _disco_s2_contraction_torch(x, self.psi.to(x.device), self.nlon_out)
# extract shape
if not _save_x_spatial_first_ok:
B, C, K, H, W = x.shape
x = x.reshape(B, self.groups, self.groupsize, K, H, W)
# do weight multiplication
out = torch.einsum("bgckxy,gock->bgoxy", x, weight_r).contiguous()
out = out.reshape(B, self.weight.shape[0], H, W)
if self.bias is not None:
out = out + self.bias.reshape(1, self.bias.shape[0], 1, 1)
return out
[docs]
class DiscreteContinuousConvTransposeS2(DiscreteContinuousConv):
r"""
Discrete-continuous (DISCO) transpose convolution on the 2-sphere, as described in :cite:`Ocampo2023`.
This is the transpose (adjoint) of
:class:`~torch_harmonics.DiscreteContinuousConvS2`. It uses the same
continuous-filter and quadrature construction but applies the
:math:`\Psi` tensor in the reverse direction -- typically to map a coarser
grid to a finer one (upsampling), analogous to a transposed/strided
convolution in the planar case. It shares the compact-support filter and
sparse, linearly scaling evaluation, and the same approximate
:math:`SO(3)` equivariance.
.. seealso::
:doc:`/guide/disco_convolutions`
User guide with the full mathematical derivation, filter basis
visualisations, and worked examples.
Parameters
----------
in_channels : int
Number of input channels
out_channels : int
Number of output channels
in_shape : Tuple[int]
Input shape of the convolution tensor
out_shape : Tuple[int]
Output shape of the convolution tensor
kernel_shape : Union[int, Tuple[int], Tuple[int, int]]
Shape of the kernel
basis_type : Optional[str]
Type of the basis functions
basis_norm_mode : Optional[str]
Mode for basis normalization
groups : Optional[int]
Number of groups
grid_in : Optional[str]
Input grid type
grid_out : Optional[str]
Output grid type
bias : Optional[bool]
Whether to use bias
theta_cutoff : Optional[float]
Theta cutoff for the filter basis functions
optimized_kernel : Optional[bool]
Whether to use the optimized kernel (if available)
References
----------
:cite:`Ocampo2023`
"""
def __init__(
self,
in_channels: int,
out_channels: int,
in_shape: Tuple[int],
out_shape: Tuple[int],
kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
basis_type: Optional[str] = "piecewise linear",
basis_norm_mode: Optional[str] = "nodal",
groups: Optional[int] = 1,
grid_in: Optional[str] = "equiangular",
grid_out: Optional[str] = "equiangular",
bias: Optional[bool] = True,
theta_cutoff: Optional[float] = None,
optimized_kernel: Optional[bool] = True,
):
super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias, optimized_kernel)
self.nlat_in, self.nlon_in = in_shape
self.nlat_out, self.nlon_out = out_shape
# make sure the p-shift works by checking that longitudes are divisible
if self.nlon_out % self.nlon_in != 0:
raise ValueError(f"nlon_out ({self.nlon_out}) must be an integer multiple of nlon_in ({self.nlon_in}) for the DISCO transpose p-shift to be exact")
# bandlimit
if theta_cutoff is None:
self.theta_cutoff = torch.pi / float(self.nlat_in - 1)
else:
self.theta_cutoff = theta_cutoff
if self.theta_cutoff <= 0.0:
raise ValueError("Error, theta_cutoff has to be positive.")
# switch in_shape and out_shape since we want the transpose convolution
idx, vals, _ = _precompute_convolution_tensor_s2(
out_shape,
in_shape,
self.filter_basis,
grid_in=grid_out,
grid_out=grid_in,
theta_cutoff=self.theta_cutoff,
transpose_normalization=True,
basis_norm_mode=basis_norm_mode,
merge_quadrature=True,
)
# sort the values
ker_idx = idx[0, ...].contiguous()
row_idx = idx[1, ...].contiguous()
col_idx = idx[2, ...].contiguous()
vals = vals.contiguous()
if self.optimized_kernel:
# preprocessed data-structure for GPU kernel
roff_idx = preprocess_psi(self.kernel_size, self.nlat_in, ker_idx, row_idx, col_idx, vals).contiguous()
self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
# save all datastructures
self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
self.register_buffer("psi_row_idx", row_idx, persistent=False)
self.register_buffer("psi_col_idx", col_idx, persistent=False)
self.register_buffer("psi_vals", vals, persistent=False)
# also store psi just in case
if not self.optimized_kernel:
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}, theta_cutoff={self.theta_cutoff}, groups={self.groups}"
@property
def psi_idx(self):
return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
[docs]
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Apply the transpose discrete-continuous convolution.
Parameters
----------
x : torch.Tensor
Input signal of shape ``(batch, in_channels, nlat_in, nlon_in)``.
Returns
-------
torch.Tensor
Convolved signal of shape ``(batch, out_channels, nlat_out, nlon_out)``.
"""
# extract shape
B, C, H, W = x.shape
x = x.reshape(B, self.groups, self.groupsize, H, W)
# do weight multiplication
x = torch.einsum("bgcxy,gock->bgokxy", x, self.weight.reshape(self.groups, self.out_per_group, self.weight.shape[1], self.weight.shape[2])).contiguous()
x = x.reshape(B, self.weight.shape[0], x.shape[-3], H, W)
if self.optimized_kernel:
out = _disco_s2_transpose_contraction_optimized(
x, self.psi_roff_idx, self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx, self.psi_vals, self.kernel_size, self.nlat_out, self.nlon_out
)
else:
out = _disco_s2_transpose_contraction_torch(x, self.psi_st.to(x.device), self.nlon_out)
if self.bias is not None:
out = out + self.bias.reshape(1, self.bias.shape[0], 1, 1)
return out
