mbtrack2.tracking.parallel module

Module to handle MPI parallel computation via the mpi4py module [1].

class Mpi(filling_pattern: ndarray[tuple[int, ...], dtype[_ScalarType_co]])

Bases: object

Class which handle parallel computation via the mpi4py module [1].

Non-empty bunches are distributed across MPI ranks as contiguous blocks: rank r owns bunches bunch_index[displacements[r]:displacements[r] + counts[r]], where bunch_index = np.where(filling_pattern)[0]. The number of ranks must not exceed the number of non-empty bunches, but any smaller rank count is allowed (a single rank can own several bunches).

Outputs of share_* methods are dicts keyed by bunch number (absolute index into filling_pattern). Every key is an element of bunch_index, so callers write beam.mpi.tau_center[bunch_num].

Parameters

filling_patternbool array of shape (h,)

Filling pattern of the beam, like Beam.filling_pattern

Attributes

commMPI.Intracomm object

MPI intra-communicator of the processor group.

rankint

Rank of the current processor.

sizeint

Number of processors in the communicator.

bunch_indexint array of shape (n_bunches,)

Non-empty bunch numbers.

countsint array of shape (size,)

Number of bunches owned by each rank.

displacementsint array of shape (size,)

Starting position (in bunch_index) of each rank’s block.

bunch_to_rank_arrint array of shape (n_bunches,)

Rank that owns each non-empty bunch.

rank_to_bunchesdict[int, np.ndarray]

Bunch numbers owned by each rank.

bunch_num_to_rankdict[int, int]

Map from bunch number to the rank that owns it.

bunch_numberslist of int

Bunch numbers owned by the current processor.

bunch_numbers_setset of int

Set version of bunch_numbers for fast membership tests.

local_countint

Number of bunches owned by the current processor.

local_displacementint

Offset of the current rank’s block inside bunch_index.

Methods

write_table(filling_pattern)

Build the rank-bunch distribution tables.

rank_to_bunch(rank)

Return the bunch numbers owned by rank.

bunch_to_rank(bunch_num)

Return the rank that owns bunch_num.

share_distributions(beam)

Compute and share bunch profiles across all ranks.

share_means(beam)

Compute and share bunch means across all ranks.

share_stds(beam)

Compute and share bunch standard deviations across all ranks.

share_scalar(beam, attr)

Gather a scalar per-bunch attribute across all ranks.

References

[1] L. Dalcin, P. Kler, R. Paz, and A. Cosimo, Parallel Distributed Computing using Python, Advances in Water Resources, 34(9):1124-1139, 2011.

__init__(filling_pattern: ndarray[tuple[int, ...], dtype[_ScalarType_co]])

write_table(filling_pattern: ndarray[tuple[int, ...], dtype[_ScalarType_co]])

Build the rank-bunch distribution tables.

Non-empty bunches are split into contiguous blocks. The first n_bunches % size ranks receive one extra bunch so the partition is as balanced as possible.

Parameters

filling_patternbool array of shape (h,)

Filling pattern of the beam, like Beam.filling_pattern

rank_to_bunch(rank: int) → ndarray[tuple[int, ...], dtype[_ScalarType_co]]

Return the bunch numbers owned by rank.

Parameters

rank : int

Returns

bunch_numbersnp.ndarray of int

Bunch numbers (absolute indices into filling_pattern) owned by rank.

bunch_to_rank(bunch_num: int) → int | None

Return the rank that owns bunch_num.

Parameters

bunch_numint

Absolute bunch number (index into filling_pattern).

Returns

rankint or None

Rank that tracks bunch_num, or None if bunch_num is empty.

_allgatherv(local: ndarray[tuple[int, ...], dtype[_ScalarType_co]], inner_size: int, dtype, mpi_dtype) → ndarray[tuple[int, ...], dtype[_ScalarType_co]]

Gather per-bunch arrays from all ranks into a single shared array.

Each rank contributes local_count rows. The result concatenates all ranks’ contributions in bunch order, producing one row per non-empty bunch aligned with self.bunch_index.

Parameters

localnp.ndarray of shape (local_count, inner_size)

Local data array for the bunches owned by the current rank.

inner_sizeint

Number of elements per bunch (columns in the result).

dtypenumpy dtype

NumPy data type for local and the returned array.

mpi_dtypeMPI.Datatype

Matching MPI datatype (e.g. MPI.DOUBLE for np.float64).

Returns

resultnp.ndarray of shape (n_bunches, inner_size)

Row j contains data for bunch self.bunch_index[j].

share_scalar(beam: Beam, attr: str) → dict[int, float]

Gather a scalar per-bunch attribute across all ranks.

Parameters

beam : Beam object attr : str

Name of the Bunch attribute to gather. Must be a scalar float.

Returns

resultdict[int, float]

Maps bunch number to the gathered attribute value.

share_distributions(beam: Beam, dimensions: str | list[str] = 'tau', n_bin: int = 75)

Compute the bunch profiles and share them across all processors.

Produces the following attributes keyed by bunch number (int):

• self.<dim>_center : dict[int, np.ndarray of shape (n_bin,)]

• self.<dim>_profile : dict[int, np.ndarray of shape (n_bin,)]

• self.<dim>_bin_length : dict[int, float]

• self.<dim>_sorted_index : dict[int, np.ndarray | None], locally-owned bunches only.

• self.charge_per_mp_all : dict[int, float]

Parameters

beam : Beam object dimensions : str or list of str, optional n_bin : int or list of int, optional

share_means(beam: Beam)

Compute and share bunch means across all processors.

Produces: - self.mean_all : dict[int, np.ndarray of shape (6,)], keyed by bunch number. - self.charge_all : dict[int, float], keyed by bunch number.

Parameters

beam : Beam object

share_stds(beam: Beam)

Compute and share bunch standard deviations across all processors.

Produces: - self.std_all : dict[int, np.ndarray of shape (6,)], keyed by bunch number. - self.charge_all : dict[int, float], keyed by bunch number.

Parameters

beam : Beam object