salvus.mesh.algorithms.distributed_mesh.scatter
salvus.mesh.algorithms.distributed_mesh.scatter salvus mesh algorithms distributed_mesh scatter Algorithms scattering objects using MPI for various data structures.
Functions
scatter_layered_model_spatially()
scatter_layered_model_spatially()def scatter_layered_model_spatially(
layered_model: LayeredModel | None,
comm: int,
chunks: (
list[tuple[float, float]] | list[tuple[float, float, float, float]]
),
root: int,
halo_width: int = 5,
) -> LayeredModel: ...Broadcast a LayeredModel from the root rank to all other ranks.
It is broadcast in a distributed fashion, with each rank receiving only the chunk of the model it is responsible for. The number of chunks must be equal to the number of ranks in the communicator.
layered_modelLayeredModel | None — The LayeredModel to broadcast. Only needs to be valid on the root rank; can be None on other ranks.commint — MPI communicator handle.chunkslist[tuple[float, float]] | list[tuple[float, float, float, float]] — List of chunks (as bounding boxes) assigned to each rank.rootint — The rank that holds the layered model to broadcast.halo_widthint — The number of points to include as a halo on each side of the chunk for discrete parameters. This prevents interpolation artifacts when the discrete parameters are divided into chunks.
scatter_xarray_spatially()
scatter_xarray_spatially()def scatter_xarray_spatially(
ds: xr.Dataset | xr.DataArray | None,
comm: int,
root: int,
chunks: (
list[tuple[float, float]] | list[tuple[float, float, float, float]]
),
coord_names: tuple[str, ...] | None = None,
halo_width: int = 5,
) -> xr.Dataset | xr.DataArray: ...Scatter an xarray Dataset or DataArray spatially across MPI ranks.
The input xarray object resides on the root rank only and after the scattering, each rank will have a local xarray object corresponding to its chunk.
The chunks are defined as ndim-1, in 2-D each chunk is x0 and x1 to denote left and right boundaries. In 3-D each chunk is x0, x1, y0, y1 to denote left, right, bottom, and top boundaries.
The model_overlap_factor defines how much overlap there is between chunks.
dsxr.Dataset | xr.DataArray | None — The xarray Dataset or DataArray to scatter. This should be provided on the root rank only.commint — The MPI communicator.rootint — The root rank where the input xarray object resides.chunkslist[tuple[float, float]] | list[tuple[float, float, float, float]] — The list of chunks defining the spatial partitioning.coord_namestuple[str, ...] | None — The names of the spatial coordinates. Defaults to (“x”,) for 2-D chunks (len=2) or (“x”, “y”) for 3-D chunks (len=4). The coordinate names map to chunk bounds as follows: - 2-D chunks (x0, x1): first coord_name is sliced by [x0, x1] - 3-D chunks (x0, x1, y0, y1): first coord sliced by [x0, x1], second coord sliced by [y0, y1]halo_widthint — The number of points to include as a halo on each side of the chunk. This prevents interpolation artifacts.
selection_with_halo_slice()
selection_with_halo_slice()def selection_with_halo_slice(
ds_coord_values: np.ndarray,
chunk_bounds: tuple[float, float],
halo_width: int,
) -> slice: ...Select all points + halo on a coordinate axis that lie within a range.
If not enough halo points are available within the dataset bounds, the selection will be trimmed. If no points lie within the chunk bounds, the nearest point to the chunk center will be selected, and the halo will be applied starting from that point instead.
ds_coord_valuesnp.ndarray — The coordinate values of the dataset.chunk_boundstuple[float, float] — A tuple of (min, max) coordinate values for the chunk.halo_widthint — The minimum number of points to include as a halo.