LinearSolvers

AMReX’s geometric multigrid solver MLMG with cell-centered, nodal and edge-centered operators, GMRES, the algebraic multigrid solver AlgMG, and a user-defined operator. These tutorials require AMReX built with AMReX_LINEAR_SOLVERS=ON (the default); CurlCurl also needs AMReX_LINEAR_SOLVERS_EM and NodalProjection AMReX_LINEAR_SOLVERS_INCFLO (both ON by default). They run in 1D, 2D and 3D, except NodalProjection (2D and 3D). Build only this topic with -DAMReXTutorials_TOPICS=LinearSolvers; see Getting Started.

All tutorials here read the MLMG options mlmg.verbose and mlmg.tol_rel (and some mlmg.max_iter) from the inputs, except GMRES, which reads verbose and tol_rel. MLMG::solve(sol, rhs, tol_rel, tol_abs) stops when the max norm of the residual is at most max(tol_abs, tol_rel*norm0), where norm0 is by default the larger of the max norms of the right-hand side and the initial residual. Set tol_abs to zero unless you know the scale of the residual.

CellPoisson

Level:: ★
Source:: LinearSolvers/CellPoisson
Key APIs:: MLPoisson, LinOpBCType, setDomainBC, setLevelBC, MLMG, MLMG::solve, setVerbose, setMaxIter

Solve \(\nabla^2\phi = f\) for a cell-centered \(\phi\) with MLPoisson. This is the workflow common to all AMReX linear solvers:

Build the operator from Geometry, BoxArray and DistributionMapping (one of each per AMR level).
Set the domain boundary types with setDomainBC, then the boundary values with setLevelBC. The values are taken from the ghost cells of a MultiFab (here the solution itself) and represent values on the domain faces.
Build an MLMG solver, set its options, and call solve.

The boundary type in each direction comes from the inputs, e.g., bc = dirichlet neumann periodic. The exact solution \(\phi = \prod_d \cos(2\pi x_d)\) works with all three. Without a Dirichlet direction, \(\phi\) is determined only up to a constant, so the program removes the mean before comparing. The problem is solved on two grids, and the error drops by a factor of four when the cell size is halved.

MLPoisson mlpoisson({geom}, {ba}, {dm});

// Domain boundary types, then the boundary values (in phi's ghost cells).
Array<LinOpBCType,AMREX_SPACEDIM> lobc, hibc;
bool has_dirichlet = false;
for (int d = 0; d < AMREX_SPACEDIM; ++d) {
    if (p.bc[d] == BCKind::dirichlet) {
        lobc[d] = hibc[d] = LinOpBCType::Dirichlet;
        has_dirichlet = true;
    } else if (p.bc[d] == BCKind::neumann) {
        lobc[d] = hibc[d] = LinOpBCType::Neumann;
    } else {
        lobc[d] = hibc[d] = LinOpBCType::Periodic;
    }
}
mlpoisson.setDomainBC(lobc, hibc);
mlpoisson.setLevelBC(0, &phi);

MLMG mlmg(mlpoisson);
mlmg.setVerbose(p.verbose);
mlmg.setMaxIter(p.max_iter);
// Stop when max|rhs - L(phi)| <= max(tol_abs, tol_rel*R), with R the
// larger of max|rhs| and the initial residual.
mlmg.solve({&phi}, {&rhs}, p.tol_rel, p.tol_abs);

Run (from the top of the repository, see Getting Started):

cmake --build build --target LinearSolvers_CellPoisson -j 8
cd build/LinearSolvers/CellPoisson
mpiexec -n 2 ./CellPoisson inputs

With mlmg.verbose = 1, MLMG prints the norms and the number of iterations:

MLMG: Initial rhs               = 116.7325837
MLMG: Initial residual (resid0) = 2821.164753
MLMG: Final Iter. 7 resid, resid/resid0 = 4.106625795e-08, 1.455649051e-11
...
n_cell = 32: max error = 0.003249084214
n_cell = 64: max error = 0.0008196247914
Convergence rate = 1.98699761

ABecLaplacian

Level:: ★★
Source:: LinearSolvers/ABecLaplacian
Key APIs:: MLABecLaplacian, setScalars, setACoeffs, setBCoeffs, LinOpBCType::inhomogNeumann, setMaxOrder, MLMG::solve with GetVecOfPtrs/GetVecOfConstPtrs

Solve the variable-coefficient equation

\[a\,\alpha(\mathbf{x})\,\phi - b\,\nabla\cdot\left(\beta(\mathbf{x})\nabla\phi\right) = f\]

with MLABecLaplacian as one composite solve over all levels of a fixed AMR hierarchy. Each finer level refines the middle half of the level below (max_level = 1 by default):

// Fixed AMR hierarchy: each finer level refines the middle half of
// the level below by a factor of two.
int const nlevels = max_level + 1;
int const ref_ratio = 2;
Vector<Geometry> geom(nlevels);
Vector<BoxArray> grids(nlevels);
Vector<DistributionMapping> dmap(nlevels);

RealBox real_box({AMREX_D_DECL(0._rt,0._rt,0._rt)}, {AMREX_D_DECL(1._rt,1._rt,1._rt)});
Array<int,AMREX_SPACEDIM> is_periodic{};
for (int d = 0; d < AMREX_SPACEDIM; ++d) { is_periodic[d] = (bc[d] == BCKind::periodic); }

Box domain(IntVect(0), IntVect(n_cell-1));
Box fine_region = domain;
for (int lev = 0; lev < nlevels; ++lev) {
    geom[lev].define(domain, real_box, CoordSys::cartesian, is_periodic);
    grids[lev].define(fine_region);
    grids[lev].maxSize(max_grid_size);
    dmap[lev].define(grids[lev]);
    fine_region.grow(-fine_region.length(0)/4).refine(ref_ratio);
    domain.refine(ref_ratio);
}

\(\alpha\) is cell-centered and \(\beta\) is face-centered. \(\beta\) changes steeply by a factor beta_in (10 by default) across a sphere covered by the fine level. The boundary type per direction is dirichlet, neumann (LinOpBCType::inhomogNeumann) or periodic. The ghost cells passed to setLevelBC hold \(\phi\) on the face for Dirichlet, and \(\partial\phi/\partial x_d\) on the face for Neumann. setMaxOrder sets the order of the stencil at Dirichlet and coarse/fine boundaries.

// The operator for all levels of the composite solve.
MLABecLaplacian mlabec(geom, grids, dmap);

// Order of the stencil at Dirichlet and coarse/fine boundaries.
mlabec.setMaxOrder(linop_maxorder);

Array<LinOpBCType,AMREX_SPACEDIM> lobc, hibc;
for (int d = 0; d < AMREX_SPACEDIM; ++d) {
    if (bc[d] == BCKind::dirichlet) {
        lobc[d] = hibc[d] = LinOpBCType::Dirichlet;
    } else if (bc[d] == BCKind::neumann) {
        lobc[d] = hibc[d] = LinOpBCType::inhomogNeumann;
    } else {
        lobc[d] = hibc[d] = LinOpBCType::Periodic;
    }
}
mlabec.setDomainBC(lobc, hibc);
for (int lev = 0; lev < nlevels; ++lev) {
    mlabec.setLevelBC(lev, &phi[lev]);
}

mlabec.setScalars(prob.a, prob.b);
for (int lev = 0; lev < nlevels; ++lev) {
    mlabec.setACoeffs(lev, acoef[lev]);
    mlabec.setBCoeffs(lev, GetArrOfConstPtrs(bcoef[lev]));
}

MLMG mlmg(mlabec);
mlmg.setVerbose(verbose);
mlmg.setMaxIter(max_iter);
mlmg.solve(GetVecOfPtrs(phi), GetVecOfConstPtrs(rhs), tol_rel, Real(0));

After the solve, MLMG has averaged the fine solution onto the covered coarse cells, and the error is checked on every level against \(\phi = \prod_d \sin(2\pi x_d + \pi/4)\).

Run (from the top of the repository, see Getting Started):

cmake --build build --target LinearSolvers_ABecLaplacian -j 8
cd build/LinearSolvers/ABecLaplacian
mpiexec -n 2 ./ABecLaplacian inputs

Try max_level=2, bc="neumann neumann neumann", or a larger n_cell.

ImplicitHeat

Level:: ★★
Source:: LinearSolvers/ImplicitHeat
Key APIs:: MLABecLaplacian, setScalars, setACoeffs/setBCoeffs with scalars, setLevelBC(0, nullptr), reusing MLMG across solves, MLMG::getNumIters

Advance the heat equation \(\partial\phi/\partial t = \kappa\nabla^2\phi\) on a periodic domain with backward Euler,

\[(1 - \Delta t\,\kappa\nabla^2)\,\phi^{n+1} = \phi^n,\]

which is MLABecLaplacian with \(a = 1\), \(\alpha = 1\), \(b = \Delta t\) and \(\beta = \kappa\). The time step is dt_factor (10 by default) times the explicit stability limit \(\Delta x^2/(2D\kappa)\); backward Euler is stable for any \(\Delta t\). Compare with the explicit HeatEquation.

The operator and the solver are built once and reused every step. Each solve starts from the previous solution:

// Build the operator and the solver once and reuse them every step.
MLABecLaplacian mlabec({geom}, {ba}, {dm});
mlabec.setDomainBC({AMREX_D_DECL(LinOpBCType::Periodic,
                                 LinOpBCType::Periodic,
                                 LinOpBCType::Periodic)},
                   {AMREX_D_DECL(LinOpBCType::Periodic,
                                 LinOpBCType::Periodic,
                                 LinOpBCType::Periodic)});
mlabec.setLevelBC(0, nullptr); // no boundary data needed for periodic
mlabec.setScalars(Real(1), dt);
mlabec.setACoeffs(0, Real(1));
mlabec.setBCoeffs(0, kappa);

MLMG mlmg(mlabec);
mlmg.setVerbose(verbose);

Real time = 0;
for (int step = 1; step <= nsteps; ++step) {
    MultiFab::Copy(phi_old, phi_new, 0, 0, 1, 0);
    // phi_new holds the previous solution, a good initial guess.
    mlmg.solve({&phi_new}, {&phi_old}, tol_rel, Real(0));
    time += dt;
    amrex::Print() << "Step " << step << ": time = " << time
                   << ", MLMG iterations = " << mlmg.getNumIters() << "\n";
}

The initial condition \(\phi = \prod_d \sin(2\pi x_d)\) is an eigenfunction of the discrete Laplacian, so each step multiplies it by a known factor, and the program checks the solution against this to the solver tolerance. It also prints the error against the PDE solution, which is dominated by the \(O(\Delta t)\) error of backward Euler.

Run (from the top of the repository, see Getting Started):

cmake --build build --target LinearSolvers_ImplicitHeat -j 8
cd build/LinearSolvers/ImplicitHeat
mpiexec -n 2 ./ImplicitHeat inputs

Try dt_factor=1 nsteps=200 to reduce the time error.

NodalProjection

Level:: ★★
Source:: LinearSolvers/NodalProjection
Key APIs:: MLNodeLaplacian, compDivergence, updateVelocity, nodal BoxArray (convert(ba, IntVect::TheNodeVector()))

Project a cell-centered velocity \(\mathbf{u}\) onto a divergence-free field, the core step of incompressible flow solvers:

\[\nabla\cdot(\sigma\nabla\phi) = \nabla\cdot\mathbf{u}, \qquad \mathbf{u}^\mathrm{new} = \mathbf{u} - \sigma\nabla\phi,\]

with \(\phi\) on the nodes, solved by MLNodeLaplacian. Here \(\sigma = 1\) is passed to the constructor; for variable density, use \(\sigma = 1/\rho\) with setSigma. The domain is periodic in \(x\) with walls (LinOpBCType::Neumann, i.e., \(\mathbf{u}\cdot\mathbf{n} = 0\)) elsewhere. The system is singular, and MLMG handles it.

compDivergence computes the nodal \(\nabla\cdot\mathbf{u}\) (the velocity needs one ghost cell), and updateVelocity subtracts \(\sigma\nabla\phi\):

// Nodal Poisson operator with constant sigma = 1.  For variable
// density, use sigma = 1/rho with setSigma instead.  The empty list
// of factories means no embedded boundaries.
Vector<FabFactory<FArrayBox> const*> const no_eb{};
MLNodeLaplacian linop({geom}, {ba}, {dm}, LPInfo(), no_eb, Real(1));
linop.setDomainBC({AMREX_D_DECL(LinOpBCType::Periodic,
                                LinOpBCType::Neumann,
                                LinOpBCType::Neumann)},
                  {AMREX_D_DECL(LinOpBCType::Periodic,
                                LinOpBCType::Neumann,
                                LinOpBCType::Neumann)});

BoxArray const nba = amrex::convert(ba, IntVect::TheNodeVector());
MultiFab rhs(nba, dm, 1, 0);
MultiFab phi(nba, dm, 1, 1);
phi.setVal(0);

// rhs = div(u) on the nodes
linop.compDivergence({&rhs}, {&vel});
Real const div0 = rhs.norm0();

// With walls and periodic boundaries only, phi is determined up to a
// constant.  MLMG handles this singular system.
MLMG mlmg(linop);
mlmg.setVerbose(verbose);
mlmg.solve({&phi}, {&rhs}, tol_rel, Real(0));

// u -= sigma*grad(phi)
linop.updateVelocity({&vel}, {&phi});

This is an approximate projection: the nodal Laplacian is not exactly the product of the discrete divergence and gradient, so \(\nabla\cdot\mathbf{u}^\mathrm{new}\) is \(O(\Delta x^2)\), not zero. The program checks that the remaining divergence and the error in the projected velocity are \(O(\Delta x^2)\).

Run (from the top of the repository, see Getting Started):

cmake --build build --target LinearSolvers_NodalProjection -j 8
cd build/LinearSolvers/NodalProjection
mpiexec -n 2 ./NodalProjection inputs

GMRES

Level:: ★★
Source:: LinearSolvers/GMRES
Key APIs:: amrex::GMRES, GMRES::define, setRestartLength, setMaxIters, getNumIters, getStatus, MLMG::precond, MLMG::setBottomSolver, GMRESMLMG

Use AMReX’s GMRES with multigrid as the preconditioner, in two ways.

Your own operator. amrex::GMRES<V,M> works with any vector type V and operator class M. ConvDiffOp implements the convection-diffusion operator \(\nabla^2 u - \mathbf{c}\cdot\nabla u\) without storing a matrix. It provides the matrix-vector product (apply), vector operations (dotProduct, linComb, …), and a preconditioner (precond) that runs one MLMG V-cycle of MLPoisson, i.e., of the Laplacian alone. The operator is not symmetric, so CG cannot be used. With the preconditioner GMRES needs about 40 iterations, independent of the grid size; without it (use_precond = 0) about 300.

void ConvDiffOp::precond (MultiFab& lhs, MultiFab const& rhs) const
{
    if (m_use_precond) {
        // Approximately solve lap(lhs) = rhs with one V-cycle, starting
        // from zero.  A fixed number of cycles keeps the preconditioner a
        // linear operator, as GMRES requires.
        lhs.setVal(0);
        m_mlmg->setPrecondIter(1);
        m_mlmg->precond({&lhs}, {&rhs}, Real(0), Real(0));
    } else {
        MultiFab::Copy(lhs, rhs, 0, 0, 1, 0);
    }
}

The preconditioner must be a fixed linear operator. That is why it does a fixed number of V-cycles and uses the smoother instead of BiCGStab as the bottom solver:

MLMG mlmg(mlpoisson);
mlmg.setVerbose(0);
// A bottom solver with a convergence test (the default BiCGStab)
// would make the preconditioner nonlinear.
mlmg.setBottomSolver(MLMG::BottomSolver::smoother);

// Part 1: GMRES for the user-defined operator.
{
    ConvDiffOp op(geom, ba, dm, vel, mlmg);
    op.usePrecond(use_precond);

    GMRES<MultiFab,ConvDiffOp> gmres;
    gmres.define(op);
    gmres.setVerbose(verbose);
    gmres.setRestartLength(30);
    gmres.setMaxIters(1000);

    MultiFab u = op.makeVecLHS();
    u.setVal(0);
    gmres.solve(u, rhs, tol_rel, Real(0));

An MLMG operator. GMRESMLMG takes an MLMG object and solves its system with GMRES, using MLMG itself as the preconditioner. This can help when MLMG alone converges slowly.

GMRESMLMG gmsolver(mlmg);
gmsolver.getGMRES().setVerbose(verbose);
gmsolver.solve(u, lap_rhs, tol_rel, Real(0));

Both problems use homogeneous Dirichlet boundaries and the exact solution \(u = \prod_d \sin(\pi x_d)\).

Run (from the top of the repository, see Getting Started):

cmake --build build --target LinearSolvers_GMRES -j 8
cd build/LinearSolvers/GMRES
mpiexec -n 2 ./GMRES inputs

Runtime parameters: n_cell, max_grid_size, velocity (the vector \(\mathbf{c}\)), use_precond, verbose and tol_rel.

AlgMG

Level:: ★★★
Source:: LinearSolvers/AlgMG
Key APIs:: AlgPartition, AlgVector, SpMatrix::define, SpMV, AlgMG, AlgMG::setKrylovSolver, MLMG::setMultigridType, BottomSolver::algmg, MLMG::setAlgMGOptions, MLMG::usedAlgMG

AlgMG is AMReX’s algebraic multigrid solver. It solves a sparse linear system \(A x = b\) and builds its coarse levels from the matrix itself, so it works where geometric multigrid struggles, e.g. for strongly varying or anisotropic coefficients. The problem here is \(-\nabla\cdot(\beta\nabla\phi) = f\) on the unit domain with \(\phi = 0\) on the boundary, where \(\beta\) is a checkerboard of blocks of block cells alternating between 1 and jump (default \(10^4\)). The right-hand side is the discrete operator applied to \(\phi = \prod_d \sin(\pi x_d)\).

Standalone. AlgPartition splits the rows, one per cell, evenly among the MPI ranks, and AlgVector is a distributed vector. The matrix is assembled from CSR arrays of the local rows. Each row has room for the \(2d+1\) entries of the stencil; a missing neighbor at the domain boundary gets column -1, which CsrValid{false} tells SpMatrix::define to drop. AlgMG then solves with V-cycles accelerated by conjugate gradients:

SpMatrix<Real> A;
A.define(partition, pv, pc, nlocal*nslots, po, CsrSorted{false}, CsrValid{false});

// The right-hand side is A times the exact solution, so the solver
// should recover it to within the tolerance.
SpMV(b, A, exact);

AlgMG<Real> amg(A); // A must outlive amg
{
    ParmParse pp("algmg");
    int verbose = 1;
    pp.query("verbose", verbose);
    amg.setVerbose(verbose);
    // V-cycles accelerated by conjugate gradients; the matrix is
    // symmetric positive definite. Also: none, bicgstab, gmres.
    AlgMGKrylovSolver krylov = AlgMGKrylovSolver::pcg;
    pp.query_enum_case_insensitive("krylov", krylov);
    amg.setKrylovSolver(krylov);
    AlgMGSmoother smoother = AlgMGSmoother::chebyshev;
    pp.query_enum_case_insensitive("smoother", smoother);
    amg.setSmoother(smoother);
}
amg.setRelTol(reltol);
amg.setThrowException(true);

x.setVal(0); // initial guess
bool converged = true;
try {
    amg.solve(x, b);

In MLMG. The same problem with MLABecLaplacian is solved with four MLMG variants:

Variant const variants[] = {
    // Geometric multigrid down to the default bottom solver
    {"geometric", MultigridType::geometric, BottomSolver::Default, false},
    // Geometric multigrid with AlgMG solving the bottom level. This does
    // not help here, because the geometric coarse levels above the
    // bottom are the problem.
    {"geometric + AlgMG bottom", MultigridType::geometric, BottomSolver::algmg, false},
    // AlgMG on the whole level in every MLMG iteration
    {"algebraic", MultigridType::algebraic, BottomSolver::Default, true},
    // Geometric first, AlgMG if the geometric iterations stall
    {"hybrid", MultigridType::hybrid, BottomSolver::Default, true}
};

Geometric multigrid diverges on this problem, because its coarse grids cannot represent the coefficient. AlgMG as the bottom solver does not help, since the geometric levels above the bottom are the problem. The algebraic type solves the whole level with AlgMG in every MLMG iteration. The hybrid type starts with geometric multigrid and, when it stalls or diverges, switches to AlgMG and restarts from the best iterate; usedAlgMG() reports the switch. setAlgMGOptions takes a callback that receives the AlgMG object MLMG builds, to change its options:

mlmg.setAlgMGOptions([] (AlgMG<Real>& amg)
{
    amg.setKrylovSolver(AlgMGKrylovSolver::pcg);
});

Use the geometric type when it converges, since it is cheaper, and the hybrid type when the coefficients may be difficult. With setThrowException(true), a failed solve throws an exception instead of aborting.

Run:

cmake --build build --target LinearSolvers_AlgMG -j 8
cd build/LinearSolvers/AlgMG
mpiexec -n 2 ./AlgMG inputs

AlgMG standalone: 7 levels, 9 iterations, relative residual 8.6e-11, max error 7.3e-10
MLMG geometric: did not converge after 88 iterations
MLMG geometric + AlgMG bottom: did not converge after 88 iterations
MLMG algebraic: 3 iterations, max error 6.6e-11
MLMG hybrid: 9 iterations, max error 9.8e-10, switched to AlgMG
PASSED

Runtime parameters: n_cell, max_grid_size, block, jump, max_iter, algmg.krylov, algmg.smoother, algmg.verbose, mlmg.verbose and mlmg.tol_rel (the tolerance of both solvers). Try jump=1, where geometric multigrid converges.

CurlCurl

Level:: ★★★
Source:: LinearSolvers/CurlCurl
Key APIs:: MLCurlCurl, setScalars, prepareRHS, MLMGT<Array<MultiFab,3>>, edge BoxArrays with amrex::convert

Solve the curl-curl equation for an edge-centered vector field,

\[\nabla\times(\alpha\nabla\times\mathbf{E}) + \beta\mathbf{E} = \mathbf{f},\]

with MLCurlCurl. Equations like this arise in implicit time stepping of Maxwell’s equations and in frequency-domain electromagnetics.

\(\mathbf{E}\) is an Array<MultiFab,3>, always with three components. \(E_d\) lives on edges along direction \(d\): cell-centered in \(d\) and nodal in the other directions. In 2D, \(E_z\) is nodal in both \(x\) and \(y\). The initial guess also supplies the tangential components on Dirichlet boundaries (here zero, a perfect conductor):

// Edge-centered fields: an Array of three MultiFabs.
Array<MultiFab,3> E, rhs, E_exact;
for (int dir = 0; dir < 3; ++dir) {
    IntVect itype(1);
    if (dir < AMREX_SPACEDIM) { itype[dir] = 0; }
    BoxArray const eba = amrex::convert(ba, itype);
    E      [dir].define(eba, dm, 1, 1);
    rhs    [dir].define(eba, dm, 1, 0);
    E_exact[dir].define(eba, dm, 1, 0);

    // Initial guess.  It also supplies the values of the tangential
    // components on Dirichlet boundaries, here zero.
    E[dir].setVal(0);

setScalars(alpha, beta) sets constant coefficients; setAlpha (nodal) and setBeta (edge-centered) set variable ones. prepareRHS makes the right-hand side consistent on edges shared by several boxes. The solver is MLMGT<Array<MultiFab,3>>, MLMG templated on the field type:

MLCurlCurl mlcc({geom}, {ba}, {dm});
mlcc.setDomainBC({AMREX_D_DECL(LinOpBCType::Dirichlet,
                               LinOpBCType::Dirichlet,
                               LinOpBCType::Dirichlet)},
                 {AMREX_D_DECL(LinOpBCType::Dirichlet,
                               LinOpBCType::Dirichlet,
                               LinOpBCType::Dirichlet)});
mlcc.setScalars(alpha, beta);

// Nodal points shared by neighboring boxes must hold the same rhs.
// Ours is computed from a formula, so this is only for illustration.
mlcc.prepareRHS({&rhs});

// MLMG templated on the field type Array<MultiFab,3>
MLMGT<Array<MultiFab,3>> mlmg(mlcc);
mlmg.setVerbose(verbose);
mlmg.solve({&E}, {&rhs}, tol_rel, Real(0));

The curl of a gradient is zero, so with a small \(\beta\) the operator is nearly singular; MLCurlCurl’s smoother is designed for this. The exact solution \(E_d = \prod_{e \ne d} \sin(\pi x_e)\) vanishes tangentially on the walls, and the error is checked to be \(O(\Delta x^2)\).

Run (from the top of the repository, see Getting Started):

cmake --build build --target LinearSolvers_CurlCurl -j 8
cd build/LinearSolvers/CurlCurl
mpiexec -n 2 ./CurlCurl inputs

Try beta=1.e-3.

CustomLinOp

Level:: ★★★
Source:: LinearSolvers/CustomLinOp
Key APIs:: MLNodeLinOp, Fapply, Fsmooth, restriction, interpolation, getNComp, buildMasks, nodalSync, setDirichletNodesToZero

Write your own linear operator and solve with MLMG. MultiCompNodeLap is a nodal operator with \(N\) coupled components,

\[L(\phi)_n = \sum_m C_{nm} \nabla^2\phi_m,\]

where \(C\) is a constant, invertible \(N \times N\) matrix from the inputs (\(N \le 4\)). It derives from MLNodeLinOp, which provides the multigrid hierarchy, ghost node exchange, the Dirichlet node mask (m_dirichlet_mask) and dot products. The derived class implements:

[[nodiscard]] int getNComp () const override { return m_ncomp; }

// out = L(in).  The ghost nodes of in have been filled.
void Fapply (int amrlev, int mglev, amrex::MultiFab& out,
             amrex::MultiFab const& in) const override;

// One sweep of block Jacobi on L(sol) = rhs.
void Fsmooth (int amrlev, int mglev, amrex::MultiFab& sol,
              amrex::MultiFab const& rhs) const override;

// Full-weighting restriction of the residual to MG level cmglev.
void restriction (int amrlev, int cmglev, amrex::MultiFab& crse,
                  amrex::MultiFab& fine) const override;

// fine += linear interpolation of the coarse correction.
void interpolation (int amrlev, int fmglev, amrex::MultiFab& fine,
                    amrex::MultiFab const& crse) const override;

void prepareForSolve () override;

restriction (full weighting) and interpolation (linear) must handle different box layouts on the two multigrid levels, which occur after agglomeration. The smoother is block Jacobi, where the diagonal block \(D = C\,\mathrm{diag}(\nabla^2)\) couples the components at a node:

// Jacobi with the diagonal blocks D = C * diag(lap):
//     sol += omega * D^{-1} (rhs - L(sol))
MultiFab Lsol(sol.boxArray(), sol.DistributionMap(), m_ncomp, 0);
Fapply(amrlev, mglev, Lsol, sol);

int const ncomp = m_ncomp;
auto const cinv = m_coeff_inv;
auto const dxinv = m_geom[amrlev][mglev].InvCellSizeArray();
Real diag_lap = 0;
for (int d = 0; d < AMREX_SPACEDIM; ++d) { diag_lap -= Real(2)*dxinv[d]*dxinv[d]; }
// Damping factor with good smoothing for the Laplacian: 2/3, 4/5, 6/7 in 1D-3D
Real const omega = Real(2*AMREX_SPACEDIM)/Real(2*AMREX_SPACEDIM+1);

auto const& sol_a = sol.arrays();
auto const& rhs_a = rhs.const_arrays();
auto const& Lsol_a = Lsol.const_arrays();
auto const& msk_a = m_dirichlet_mask[amrlev][mglev]->const_arrays();

ParallelFor(sol, [=] AMREX_GPU_DEVICE (int b, int i, int j, int k)
{
    if (msk_a[b](i,j,k)) { return; }
    GpuArray<Real,max_ncomp> res{};
    for (int m = 0; m < ncomp; ++m) {
        res[m] = rhs_a[b](i,j,k,m) - Lsol_a[b](i,j,k,m);
    }
    for (int n = 0; n < ncomp; ++n) {
        Real z = 0;
        for (int m = 0; m < ncomp; ++m) { z += cinv[n*ncomp+m] * res[m]; }
        sol_a[b](i,j,k,n) += omega * z / diag_lap;
    }
});

// Make nodes shared by several boxes consistent.
nodalSync(amrlev, mglev, sol);

The operator is single-level only. A multi-level operator would also implement averageDownSolutionRHS and reflux for the coarse/fine boundaries (see Multi-Component Operators in the AMReX documentation). The program checks the error against the exact solution \(\phi_n = \prod_d \sin((n+1)\pi x_d)\), which is zero on the Dirichlet boundary.

Run (from the top of the repository, see Getting Started):

cmake --build build --target LinearSolvers_CustomLinOp -j 8
cd build/LinearSolvers/CustomLinOp
mpiexec -n 2 ./CustomLinOp inputs

Try three components: coeff="1 0 0 0 2 0.5 0 0.3 1".