S2 Partitioning

The S2Partition class will create a equal-area decomposition of a spherical shell using the number of available MPI ranks. After initialization, the instance contains information about the decomposition and the local rank coordinates.

The decomposition follows the “Recursive Zonal Equal Area Sphere Partitioning” algorithm by Paul Leopardi. The spherical shell will be divided into two polar caps (first and last MPI rank), and a number of rings with a variable of segments in between. All “cuts” are at constant phi or theta, which simplifies the implementation of ghost zones between neighbors.

from mpipartition import S2Partition

# partitioning S2 among the available ranks
partition = Partition()

# print theta and phi extent of all ranks:
print(
    f"Rank {partition.rank}:\n"
    f"  theta: [{partition.theta_extent[0]:5.3f}, {partition.theta_extent[1]:5.3f}]\n"
    f"  phi  : [{partition.phi_extent[0]:5.3f}, {partition.phi_extent[1]:5.3f}]\n"
    f"  area : {partition.area:5.3f}"
)

# print size of this rank (as fraction of unit-cube).
# Note: the extent of each rank will be the same
if partition.rank == 0:
    print(partition.extent)

Example Partitions

You can use the mpipartition-s2 executable to obtain the decomposition information for a given number of ranks (and visualize the decomposition). Here are some examples for 5, 10 and 100 ranks.

5 MPI ranks

S2 partitioning with 5 ranks — Segmentation of S2 with 5 MPI ranks.

Segmentation statistics for 5 ranks:
polar cap angle: 0.927295
number of rings: 1
   ring   0:   3 segments between theta=[0.927295, 2.214297]]
Segment area imbalance:
   max/min: 1.000000
   max/avg: 1.000000
Total edge/area ratio: 2.214498

10 MPI ranks

S2 partitioning with 10 ranks — Segmentation of S2 with 10 MPI ranks.

Segmentation statistics for 10 ranks:
polar cap angle: 0.643501
number of rings: 2
   ring   0:   4 segments between theta=[0.643501, 1.570796]]
   ring   1:   4 segments between theta=[1.570796, 2.498092]]
Segment area imbalance:
   max/min: 1.000000
   max/avg: 1.000000
Total edge/area ratio: 3.380669

100 MPI ranks

S2 partitioning with 100 ranks — Segmentation of S2 with 100 MPI ranks.

Segmentation statistics for 100 ranks:
polar cap angle: 0.200335
number of rings: 8
   ring   0:   6 segments between theta=[0.200335, 0.535527]]
   ring   1:  11 segments between theta=[0.535527, 0.876298]]
   ring   2:  15 segments between theta=[0.876298, 1.223879]]
   ring   3:  17 segments between theta=[1.223879, 1.570796]]
   ring   4:  17 segments between theta=[1.570796, 1.917713]]
   ring   5:  15 segments between theta=[1.917713, 2.265295]]
   ring   6:  11 segments between theta=[2.265295, 2.606066]]
   ring   7:   6 segments between theta=[2.606066, 2.941258]]
Segment area imbalance:
   max/min: 1.000000
   max/avg: 1.000000
Total edge/area ratio: 11.206372

S2 Distribution Algorithms

Processing large datasets on multiple MPI ranks requires to distribute the data among the processes. The mpipartition package contains the following functions for data on the sphere:

`s2_distribute`	Distribute particles among MPI ranks according to the S2 partition.
`s2_overload`	Copy data within an overload angle to the neighboring ranks ("ghost" particles)

Distribution/Overload Examples

In the following example, we generate 100 randomly positioned points per rank and then distribute them according to the angular coordinates.

from mpipartition import S2Partition
from mpipartition import s2_distribute, s2_overload

# decompose a sphere with the available MPI ranks (equal area)
partition = S2Partition()

# number of random particles per rank
n_local = 100

# randomly distributed particles in a cube spanning [-1, 1]^3
data = {
    "x": np.random.uniform(-1, 1, n_local),
    "y": np.random.uniform(-1, 1, n_local),
    "z": np.random.uniform(-1, 1, n_local),
    "id": n_local * partition.rank + np.arange(n_local),
    "rank": np.ones(n_local) * partition.rank
}

# calculate angular coordinates
data['theta'] = np.arccos(data['z'])
data['phi'] = np.arctan2(data['y'], data['x']) + np.pi

# assign to rank by position
data_distributed = s2_distribute(partition, data)

Now, we overload the partitions by 0.1 radians:

data_overloaded = s2_overload(partition, data_distributed, 0.1)

References

S2Partition

class mpipartition.S2Partition(*, equal_area=True, comm=None, verbose=False)[source]

An MPI decomposition of the spherical shell into equal-area segments

Parameters:

equal_area (bool) – If True, the spherical shell is divided into equal-area segments. If False, use equally spaced rings (in theta) instead.
comm (MPI.Comm) – The MPI communicator to use for the decomposition (default: COMM_WORLD)
verbose (bool) – If True, rank 0 will print information about the segmentation.

property all_phi_extents: ndarray[Any, dtype[float64]]: the phi extent of all segments, shape (nranks, 2)

property all_theta_extents: ndarray[Any, dtype[float64]]: the theta extent of all segments, shape (nranks, 2)

property area: float: the area of the segment assigned to this rank

property comm: MPI Communicator

property equal_area: bool: whether the partition is equal-area

property nranks: int: the total number of processors

property phi_extent: Tuple[float, float]: the phi extent of the segment assigned to this rank

property rank: int: the MPI rank of this processor

property ring_dtheta: float | None: the theta spacing between rings (only for non-equal-area partitions)

property ring_segments: ndarray[Any, dtype[int64]]: the number of segments in each ring

property ring_thetas: ndarray[Any, dtype[float64]]: the theta boundaries of all rings

property theta_cap: float: the polar cap angle

property theta_extent: Tuple[float, float]: the theta extent of the segment assigned to this rank

s2_distribute

mpipartition.s2_distribute(partition, data, *, theta_key='theta', phi_key='phi', verbose=False, verify_count=True, validate_home=False)[source]

Distribute particles among MPI ranks according to the S2 partition.

Parameters:

partition (S2Partition) – The S2 partition to use for the distribution.
data (Mapping[str, ndarray]) – The particle data to distribute, as a collection of 1-dimensional arrays. Each array must have the same length (number of particles) and the map needs to contain at least the keys theta_key and phi_key.
theta_key (str) – The key in data that contains the particle theta coordinates (latitude), in the range [0, pi].
phi_key (str) – The key in data that contains the particle phi coordinates (longitude), in the range [0, 2*pi].
verbose (bool | int) – If True, print summary statistics of the distribute. If > 1, print statistics of each rank (i.e. how much data each rank sends to every other rank).
verify_count (bool) – If True, make sure that total number of objects is conserved.
validate_home (bool) – If True, validate that each rank indeed owns the particles that it was sent.

Returns:

data – The distributed particle data (i.e. the data that this rank owns)

Return type:

ParticleDataT

s2_overload

mpipartition.s2_overload(partition, data, overload_angle, *, theta_key='theta', phi_key='phi', verbose=False)[source]

Copy data within an overload angle to the neighboring ranks (“ghost” particles)

This method assumes that the particle data is already correctly distributed, i.e. that all particles on a given rank are within the bounds of the rank’s segment.

Parameters:

partition (S2Partition) – The MPI partition defining which rank should own which subvolume of the data
data (Mapping[str, ndarray]) – The particle data to be redistributed, as a collection of 1-dimensional arrays. Each array must have the same length (number of particles) and the map needs to contain at least the keys theta_key and phi_key.
overload_angle (float) – The overload angle in radians. Particles within this angle of a rank’s segment will be copied to the neighboring ranks. Note that the angle can be at maximum half of the smallest segment size in the partition.
theta_key (str) – The key in data that contains the particle theta coordinates (latitude), in the range [0, pi].
phi_key (str) – The key in data that contains the particle phi coordinates (longitude), in the range [0, 2*pi].
verbose (bool | int) – If True, print summary statistics of the distribute. If > 1, print statistics of each rank (i.e. how much data each rank sends to every other rank).

Returns:

data – The combined data of objects within the rank’s segmentt as well as the objects within the overload angle of neighboring ranks.

Return type:

ParticleDataT

Notes

The function does not change the objects’ coordinates or alter any data. Objects that have been overloaded accross the periodic boundary at 0 and 2pi will still have the original phi. In case “local” coordinates are required, this will need to be done manually after calling this function.

S2 decomposition visualization

mpipartition-s2

mpipartition-s2 [OPTIONS] NRANKS

Options

--equal-area: If set, partition S2 into equal area segments by varying delta_theta of rings.

--precision <precision>: Number of decimal places to use for printing.

--figure <figure>: If set, save a visual representation of the S2 partitioning to this file.

--use-mollweide: If set, use the Mollweide projection. Otherwise, use a rectangular theta-phi plot for visualization.

--figure-pad <figure_pad>: pad_inches argument used for saving the figure. Set to 0 to remove whitespace.

Arguments

NRANKS: Required argument