torch_harmonics.distributed.compute_polar_halo_radius#

torch_harmonics.distributed.compute_polar_halo_radius(
lats_in,
lats_out,
theta_cutoff,
lat_in_shapes,
lat_out_shapes,
)[source]#

Latitude halo radius a polar shard needs from its neighbours, derived analytically.

A localized operator couples an output latitude only to input latitudes within theta_cutoff, a contiguous band given by latitude_support_band(). A polar rank owns a contiguous range of both axes, so the rows it must borrow are exactly those where the band of its output range runs past its own input range. This returns the largest such overhang over all ranks, which is the halo width polar_halo_exchange() has to move.

The band bounds are monotone in the output latitude, so only the two ends of each rank’s output range are examined, and the whole thing costs a pair of binary searches rather than a pass over a sparsity pattern. Computing it from the geometry rather than by measuring a precomputed pattern is also what lets the pattern itself be built on the band – otherwise the radius would depend on the very tensor it is supposed to bound.

It is evaluated identically on every rank, so no communication is involved.

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) – Effective angular support radius, i.e. including the widening the sparsity pattern is built with, see effective_theta_cutoff().

  • lat_in_shapes (List[int]) – Input latitudes held by each polar rank, ordered by rank.

  • lat_out_shapes (List[int]) – Output latitudes held by each polar rank, ordered by rank.

Returns:

Halo radius in latitude rows; 0 when the polar group holds a single rank.

Return type:

int

Raises:

ValueError – If the radius exceeds the smallest local input chunk. polar_halo_exchange() only communicates with immediate neighbours, so a halo wider than a neighbour’s entire share would have to reach past it to the next rank out.