MPI¤
ExchangePlan(
layout: _LocalLayout,
local_sparsity_pattern: csr_matrix | None = None,
*,
comm: Comm,
)
An MPI communication plan for parallel FEM assembly based on point-to-point exchanges of ghost DOF values.
Parameters:
-
(layout¤_LocalLayout) –local DOF layout information for this rank
-
(local_sparsity_pattern¤csr_matrix | None, default:None) –optional local sparsity pattern for Hessian assembly.
-
(comm¤Comm) –MPI communicator
Methods:
-
make_scatter_fwd_set–Return a JIT'd function: x_owned → u_local.
-
make_scatter_rev_add–Return a JIT'd function: u_local → owned_data.
Attributes:
-
layout– -
hessian_layout(_HessianLayout | None) – -
global_size(int) –Total number of free DOFs in the global linear system.
-
rstart(int) –First owned DOF index.
-
rend(int) –One-past-last owned DOF index.
-
local_size(int) –Number of DOFs owned by this rank.
-
owned_nnz(int) –Number of sparse matrix entries owned by this rank.
-
owned_csr(tuple[NDArray[int32], NDArray[int32]]) –(indptr, indices) for owned rows.
property
¤
(indptr, indices) for owned rows.
Return a JIT'd function: x_owned → u_local.
Gathers ghost DOF values from owning ranks.
Return a JIT'd function: u_local → owned_data.
Computes the local function, then gathers contributions from ghosts. If
is_hessian=True, the local function must return a ColoredMatrix.
AllreducePlan(
global_size: int,
global_sparsity_pattern: csr_matrix | None = None,
*,
comm: Comm,
)
MPI communication plan for allreduce-based parallel FEM assembly.
Every rank holds the full replicated solution vector. Each rank computes local gradient/hessian contributions over its element subset; an allreduce sums them into the global result on every rank. Each rank then assembles only its owned rows into the distributed solver.
Use when DOF coupling is non-nearest-neighbor (cohesive interfaces, contact) or as a simpler starting point before switching to ExchangePlan for better weak scaling.
Parameters:
-
(global_size¤int) –total number of free DOFs in the global linear system
-
(global_sparsity_pattern¤csr_matrix | None, default:None) –reduced sprasity pattern of the global Hessian (only needed for Hessian assembly; pass None for gradient-only use)
-
(comm¤Comm) –MPI communicator
Methods:
-
make_allgather–Return a function: x_owned → u_free_global via Allgatherv.
-
make_allreduce_owned–Return a JIT'd function that allreduces local_fn output → owned rows.
Attributes:
-
global_size(int) –Total number of free DOFs in the global linear system.
-
rstart(int) –First owned DOF index.
-
rend(int) –One-past-last owned DOF index.
-
local_size(int) –Number of DOFs owned by this rank under the block distribution.
-
owned_nnz(int) –Number of sparse matrix entries owned by this rank.
-
owned_csr(tuple[ndarray, ndarray]) –(indptr, indices) for owned rows — pass directly to solver preallocation.
property
¤
Number of DOFs owned by this rank under the block distribution.
Pass this as the explicit local size when creating distributed solver objects so their row distribution matches this plan.
property
¤
(indptr, indices) for owned rows — pass directly to solver preallocation.
Return a function: x_owned → u_free_global via Allgatherv.
Gathers the owned DOF slice from every rank into the full replicated solution vector needed for local FEM computation on each rank.
NOT JIT-compatible: uses MPI.Allgatherv directly. Call this function in the PETSc callback before passing its result into any JIT-compiled code.
Return a JIT'd function that allreduces local_fn output → owned rows.
Computes the local function (gradient or hessian) over the full replicated DOF vector, allreduces across all ranks, then returns only the owned slice.
Parameters:
-
–local_fn¤function returning a vector (gradient) or ColoredMatrix (hessian)
-
–is_hessian¤whether the local_fn returns a ColoredMatrix (True) or a dense vector (False).