Python

These tutorials use pyAMReX, the Python bindings of AMReX, with zero-copy NumPy (CPU), CuPy (CUDA/HIP) or dpnp (SYCL) views of AMReX data. They are Python scripts written for 3D (amrex.space3d) and need pyAMReX 26.09 or newer (import amrex), plus CuPy for CUDA/HIP or dpnp for SYCL GPU builds of pyAMReX; see the pyAMReX documentation for installation. As with the C++ programs, arguments after the script name are passed to AMReX, so python3 main.py inputs n_cell=64 reads the inputs file and overrides n_cell. The scripts are not compiled; with -DAMReXTutorials_PYTHON=ON (and -DPython3_EXECUTABLE pointing to a Python that has pyAMReX), CTest runs them as the tests Python_<Name>. See Getting Started. The MPMD tutorial also has a pyAMReX program.

MultiFab

Level:

★

Source:

Python/MultiFab

Key APIs:

amr.initialize, amr.Box, amr.BoxArray, amr.DistributionMapping, amr.Geometry, amr.MultiFab, amr.xp, mf.array(mfi).to_xp(), mf.to_xp(), mf.fill_boundary, mf.min, mf.max, mf.sum, amr.write_single_level_plotfile, amr.ParmParse

The Python version of MultiFab: build a MultiFab, fill it through zero-copy views, fill ghost cells, compute reductions, and write a plotfile.

for mfi in mf: loops over the boxes owned by this MPI rank, and mf.array(mfi).to_xp(copy=False) returns a NumPy, CuPy or dpnp array that shares memory with the box’s data. The view includes the ghost cells, is indexed [x, y, z, component] with indices starting at 0, so the valid cells are [ng:ng+nx, ng:ng+ny, ng:ng+nz] for a box of nx*ny*nz cells. mf.to_xp() returns the views of all local boxes as a list. Pass order="C" for [component, z, y, x] indexing, which is usually faster on GPUs.

# Loop over the boxes owned by this rank. mf.array(mfi).to_xp() is a
# zero-copy view of the box's data including ghost cells, indexed
# [x, y, z, component] with local indices starting at 0.
for mfi in mf:
    bx = mfi.validbox()
    phi = mf.array(mfi).to_xp(copy=False)
    x = (xp.arange(bx.small_end[0], bx.big_end[0] + 1) + 0.5) * dx[0]
    y = (xp.arange(bx.small_end[1], bx.big_end[1] + 1) + 0.5) * dx[1]
    z = (xp.arange(bx.small_end[2], bx.big_end[2] + 1) + 0.5) * dx[2]
    r2 = (
        (x[:, None, None] - 0.5) ** 2
        + (y[None, :, None] - 0.5) ** 2
        + (z[None, None, :] - 0.5) ** 2
    )
    # Valid cells only: skip ng ghost cells on each side
    valid = phi[ng : ng + x.size, ng : ng + y.size, ng : ng + z.size, 0]
    valid[...] = 1.0 + xp.exp(-r2 / 0.01)

# All local boxes at once: a list of views, one per box
for phi in mf.to_xp():
    phi[..., 1] = 2.0 * phi[..., 0]

Command-line arguments, such as an inputs file, are passed to AMReX, and all AMReX objects must be destroyed before amr.finalize:

if __name__ == "__main__":
    if not hasattr(amr, "xp"):
        sys.exit("This tutorial needs pyAMReX 26.09 or newer.")
    # Command line arguments, e.g., an inputs file, are passed to AMReX.
    amr.initialize(sys.argv[1:])
    status = main()  # AMReX objects must be gone before finalize
    amr.Print("PASSED" if status == 0 else "FAILED")
    amr.finalize()
    sys.exit(status)

Run:

cd Python/MultiFab
python3 main.py inputs
mpiexec -n 2 python3 main.py inputs n_cell=64

min including ghost cells before FillBoundary: 0.0
min including ghost cells after  FillBoundary: 1.0
max(phi)      = 1.9293757195853387
integral(phi) = 1.0055683279968117 (exact: 1.0055683279968317)
max(2 phi)    = 3.8587514391706774
Wrote plt00000
PASSED

HeatEquation

Level:

★

Source:

Python/HeatEquation

Key APIs:

amr.MultiFab, amr.xp, mf.array(mfi).to_xp(), mf.fill_boundary, geom.periodicity(), mf.max, amr.ParmParse.to_dict, amr.write_single_level_plotfile

The Python version of HeatEquation: the heat equation \(\partial\phi/\partial t = D\nabla^2\phi\) with explicit time steps on a periodic unit cube. The initial condition is \(1 + \sin(2\pi x)\sin(2\pi y)\sin(2\pi z)\); its sine mode decays as \(\exp(-12\pi^2 D t)\), and the final error against this exact solution is checked.

The stencil is written with array slicing on the zero-copy view of each box. The loop over boxes is in Python, but the work on each box is done by NumPy, CuPy or dpnp, so large boxes (max_grid_size) keep the Python overhead small.

# Fill ghost cells, here all from periodic images and neighbors.
phi_old.fill_boundary(geom.periodicity())

# phi_new = phi_old + dt * D * Laplacian(phi_old), with array slicing
fac = dt_step * diffusion_coef
for mfi in phi_new:
    po = phi_old.array(mfi).to_xp(copy=False)[..., 0]
    pn = phi_new.array(mfi).to_xp(copy=False)[..., 0]
    c = po[v]
    pn[v] = c + fac * (
        (po[2:, 1:-1, 1:-1] - 2.0 * c + po[:-2, 1:-1, 1:-1]) / dx[0] ** 2
        + (po[1:-1, 2:, 1:-1] - 2.0 * c + po[1:-1, :-2, 1:-1]) / dx[1] ** 2
        + (po[1:-1, 1:-1, 2:] - 2.0 * c + po[1:-1, 1:-1, :-2]) / dx[2] ** 2
    )

Run:

cd Python/HeatEquation
python3 main.py inputs
mpiexec -n 2 python3 main.py inputs n_cell=64 plot_int=20

Step 10, time = 0.0014648437499999998
...
Finished 69 steps at time 0.01, mode amplitude = 0.3059442057
Max error = 0.001952850866, tolerance = 0.008614567486
PASSED