torch_harmonics.quadrature.latitude_support_band#

torch_harmonics.quadrature.latitude_support_band(lats_in, lats_out, theta_cutoff)[source]#

Inclusive range of input latitudes that can lie within theta_cutoff of each output latitude.

The localized operators on the sphere – the DISCO convolutions and neighborhood attention – only couple points closer than an angular cutoff. A colatitude difference is bounded by the great-circle distance between the two points, so

\[|\theta_\mathrm{in} - \theta_\mathrm{out}| \le \theta_c\]

is a necessary condition for a pair to interact, and the latitudes satisfying it form a contiguous band around each output latitude. Everything outside the band is known to be zero without evaluating it.

That band is what makes these operators cheap to build and to distribute. The sparsity pattern only has to be evaluated inside it, which is the difference between \(O(N_\theta^2 N_\lambda)\) and \(O(N_\theta N_\lambda b)\) work; and a polar shard only ever needs latitudes inside it from its neighbours, which is what bounds the halo.

The band is a superset of the true support, since it ignores longitude and so admits points that are close in latitude but far apart on the sphere. It is tight in the sense that it cannot be narrowed without knowing the longitudes: a point at the output’s own longitude attains the bound.

Because the criterion is symmetric in the two grids, it does not matter which of them the caller regards as the input – the transpose direction may pass them the other way round.

Parameters:

  • lats_in (torch.Tensor) – Input colatitudes in radians, ascending, shape (nlat_in,).

  • lats_out (torch.Tensor) – Output colatitudes in radians, ascending, shape (nlat_out,).

  • theta_cutoff (float) – Angular support radius of the filter basis. Pass the same effective value the sparsity pattern is built with, including any widening factor, or the band may exclude entries the pattern would keep.

Returns:

  • lo (torch.Tensor) – First input-latitude index in the band, per output latitude, shape (nlat_out,).

  • hi (torch.Tensor) – Last input-latitude index in the band, inclusive, shape (nlat_out,). hi < lo marks an output latitude that no input latitude can reach.

Return type:

Tuple[Tensor, Tensor]

Notes

The bound is compared inclusively, but only up to floating point: a node sitting exactly theta_cutoff away falls on whichever side the arithmetic rounds to, and that can differ between platforms. It does not matter for the cutoffs in use, which come from the grid’s node spacing widened by THETA_CUTOFF_EPS and so land between nodes rather than on one. A cutoff chosen to coincide exactly with a node separation is the case to avoid.

Examples

>>> import torch
>>> from torch_harmonics.quadrature import latitude_support_band, precompute_latitudes
>>> lats, _ = precompute_latitudes(16)
>>> cutoff = 1.5 * float(lats[1] - lats[0])   # one and a half grid spacings
>>> lo, hi = latitude_support_band(lats, lats, cutoff)
>>> int(lo[8]), int(hi[8])
(7, 9)
>>> int(lo[0]), int(hi[0])                    # clamped at the pole
(0, 1)