zoomy_core.symsolver.numpy_backend.parallel

zoomy_core.symsolver.numpy_backend.parallel#

Distributed runs of the NumPy backend (mpi4py): the halo exchange and the global reductions of the runtime’s name table, the global sums and norms of the implicit solves, and the distributed linear solves. The partition – owned cells, two halo layers, the exchange lists, the cells and faces each rank counts – is zoomy_core.mesh.distributed.Partition, the JAX backend’s.

mpi4py and petsc4py are imported only when a run is distributed: the serial backend (and the browser) runs without them.

Linear solves (MpiLinearSolver): PETSc through petsc4py – KSP GMRES on the matrix-free operator, FGMRES + BoomerAMG (or GAMG) on the colour-assembled rows, a parallel LU (MUMPS, SuperLU_dist) of the implicit diffusion’s Jacobian – or, without petsc4py, a distributed restarted FGMRES with block-Jacobi preconditioners from the rank’s own block.

References

Module Contents#

Classes#

MpiLinearSolver

The linear solves of a distributed NumPy run (pressure, implicit diffusion).

no_halo

Context manager inside which halo is the identity: an operator then gives the rank’s own rows with its halo columns as probed (assembly by colouring, no collective inside).

Distributed

One rank’s piece of a distributed NumPy run: halo exchange, global reductions and sums (each cell and face counted on one rank), the global numbering of the owned unknowns and the distributed linear solves.

Functions#

shell_gmres

GMRES of the matrix-free mv (owned -> owned) for the owned b: PETSc’s KSP on a shell matrix, or the distributed FGMRES – the serial backend’s unpreconditioned GMRES, with global inner products.

amg_gmres

GMRES preconditioned by AMG on the colour-assembled owned rows, left-scaled by the inverse cell blocks (as assembly.host_amg_solve): PETSc FGMRES + BoomerAMG / GAMG, the hierarchy lagged (amg_rebuild, amg_growth); or (krylov) the distributed FGMRES with the rank’s own smoothed-aggregation V-cycle (block Jacobi).

Data#

PREFIX

API#

class zoomy_core.symsolver.numpy_backend.parallel.MpiLinearSolver#

Bases: zoomy_core.misc.options.Option

The linear solves of a distributed NumPy run (pressure, implicit diffusion).

Initialization

Initialize self. See help(type(self)) for accurate signature.

auto#

‘auto’

petsc#

‘petsc’

krylov#

‘krylov’

class zoomy_core.symsolver.numpy_backend.parallel.no_halo#

Context manager inside which halo is the identity: an operator then gives the rank’s own rows with its halo columns as probed (assembly by colouring, no collective inside).

zoomy_core.symsolver.numpy_backend.parallel.PREFIX#

‘zoomy_numpy_’

class zoomy_core.symsolver.numpy_backend.parallel.Distributed(part, mesh_np, method=None)#

One rank’s piece of a distributed NumPy run: halo exchange, global reductions and sums (each cell and face counted on one rank), the global numbering of the owned unknowns and the distributed linear solves.

Unknowns of a (k, cells) block are numbered cell-major: local cell c, row r -> c k + r; the owned ones globally gid(c) k + r, a contiguous range per rank (Partition.offsets).

Initialization

halo(X, mesh=None)#

X with its halo columns refreshed from their owners (point-to-point); the identity on a scalar and inside no_halo.

owned(x, empty)#

The entries of x (faces or cells on the last axis) this rank counts; the others empty.

gather(x)#

Every rank’s x, in rank order (the same list on every rank).

reduce_min(x)#

Global min over the counted entries; a NaN anywhere is the result (rank-local np.min, then the min of the gathered values: no MPI.MIN comparison drops it).

reduce_max(x)#

Global max over the counted entries; NaN-propagating as reduce_min.

reduce_any(x)#

Global any over the counted entries.

total(x)#

Sum over the ranks of x (scalar or array), in rank order: the same bits on every rank, so decisions taken on it agree.

norm(r)#

Euclidean norm of r (..., cells) over the owned cells of all ranks.

integral(X, volumes)#

sum_cells X vol per row over the whole mesh (e.g. the mass), on every rank.

blocks()#

The name-table entries of a distributed run.

own(X)#

The owned unknowns of X (k, cells), cell-major.

embed(xo, k)#

(k, cells) with the owned unknowns xo and zero halo columns.

assemble(op, k, nbrs, key)#

The owned rows of the linear operator op (k, cells) -> (k, cells) on the cell stencil nbrs, by colouring with the halo off (no collective inside op, so the ranks may take different numbers of colours).

Returns

(rows, cols, vals): CSR-ordered owned rows, local cell-major columns.

global_csr(rows, cols, vals, k)#

(indptr, global columns, values) of owned rows rows, columns sorted per row.

log_once(key, msg)#

Log msg on rank 0, once per key.

krylov(mv, b, tol, restart, maxiter, precond=None)#

Restarted FGMRES of mv(x) = b on the owned unknowns (right preconditioner precond), global inner products (CGS2, two gathers per iteration); stops at |r| <= tol |b| or after restart x maxiter iterations.

References

local_block(rows, cols, vals, k)#

The rank’s rows as a scipy CSR (k n_own, k cells) (local columns).

apply(J, xo, k)#

J (local block) times the owned unknowns xo, their halo copies exchanged.

lu(rows, cols, vals, k)#

A solver solve(rhs_owned) -> x_owned of the assembled rows: PETSc’s parallel LU (MUMPS, else SuperLU_dist; else GMRES to 1e-14), or (krylov) FGMRES to 1e-14 preconditioned by the LU of the rank’s diagonal block. close() frees it.

zoomy_core.symsolver.numpy_backend.parallel.shell_gmres(dist, mv, b, tol, restart, maxiter)#

GMRES of the matrix-free mv (owned -> owned) for the owned b: PETSc’s KSP on a shell matrix, or the distributed FGMRES – the serial backend’s unpreconditioned GMRES, with global inner products.

zoomy_core.symsolver.numpy_backend.parallel.amg_gmres(dist, op, b, k, nbrs, ls, cache)#

GMRES preconditioned by AMG on the colour-assembled owned rows, left-scaled by the inverse cell blocks (as assembly.host_amg_solve): PETSc FGMRES + BoomerAMG / GAMG, the hierarchy lagged (amg_rebuild, amg_growth); or (krylov) the distributed FGMRES with the rank’s own smoothed-aggregation V-cycle (block Jacobi).

Args

  • op: the linear operator (k, cells) -> (k, cells) (halo inside)

  • b: owned right-hand side

  • cache: kept between the solves of one operator

Contents