From 4d3cb9781efa79f805dd5d46fd7a6aac65aa27ae Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Fri, 5 Jun 2026 13:11:21 -0700 Subject: [PATCH] Modernize Python tutorials - HeatEquation: read runtime parameters via ParmParse from the C++ inputs file (closes the long-standing TODO); forward command line arguments to AMReX so `python3 main.py ../Exec/inputs` mirrors the C++ binary; use the `Array4.to_xp()` zero-copy views. - MultiFab: use `Array4.to_xp()`, forward command line arguments, demonstrate global min/max/sum reductions. - MPMD Case-2: forward command line arguments. - CI: pass the inputs file to the HeatEquation Python runs and add an MPI-parallel (2 ranks) Python run. - Docs: rewrite the Python tutorials page with install/run conventions. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/linux.yml | 21 ++++---- Docs/source/Python_Tutorial.rst | 39 ++++++++++++--- ExampleCodes/MPMD/Case-2/main.py | 6 ++- GuidedTutorials/HeatEquation/Source/main.py | 54 ++++++++++++++------- GuidedTutorials/MultiFab/main.py | 19 +++++--- 5 files changed, 97 insertions(+), 42 deletions(-) diff --git a/.github/workflows/linux.yml b/.github/workflows/linux.yml index 25a58372..7ffff81a 100644 --- a/.github/workflows/linux.yml +++ b/.github/workflows/linux.yml @@ -38,9 +38,10 @@ jobs: - name: Run Python run: | cd GuidedTutorials/MultiFab/ - python main.py + python3 main.py cd ../HeatEquation/Source/ - python main.py + python3 main.py ../Exec/inputs + mpiexec -n 2 python3 main.py ../Exec/inputs # Build all tutorials tutorials_omp: @@ -69,9 +70,9 @@ jobs: - name: Run Python run: | cd GuidedTutorials/MultiFab/ - python main.py + python3 main.py cd ../HeatEquation/Source/ - python main.py + python3 main.py ../Exec/inputs tutorials_clang: name: Clang SP Particles DP Mesh Debug [tutorials] @@ -104,9 +105,9 @@ jobs: - name: Run Python run: | cd GuidedTutorials/MultiFab/ - python main.py + python3 main.py cd ../HeatEquation/Source/ - python main.py + python3 main.py ../Exec/inputs # Build all tutorials w/o MPI tutorials-nonmpi: @@ -135,9 +136,9 @@ jobs: - name: Run Python run: | cd GuidedTutorials/MultiFab/ - python main.py + python3 main.py cd ../HeatEquation/Source/ - python main.py + python3 main.py ../Exec/inputs # Build all tutorials tutorials-nofortran: @@ -164,9 +165,9 @@ jobs: - name: Run Python run: | cd GuidedTutorials/MultiFab/ - python main.py + python3 main.py cd ../HeatEquation/Source/ - python main.py + python3 main.py ../Exec/inputs # Build all tutorials with CUDA tutorials-cuda: diff --git a/Docs/source/Python_Tutorial.rst b/Docs/source/Python_Tutorial.rst index 8ae96a8f..af8fe472 100644 --- a/Docs/source/Python_Tutorial.rst +++ b/Docs/source/Python_Tutorial.rst @@ -4,17 +4,44 @@ Python ====== -These examples show how to use AMReX from Python. +These examples show how to use AMReX from Python, via `pyAMReX `__. AMReX applications can also be interfaced to Python with the same logic. +Installation +============ + In order to run the Python tutorials, you need to have pyAMReX installed. -Please see `pyAMReX `__ for more details. +Please see the `pyAMReX documentation `__ for installation details. Alternatively, you can build the ExampleCodes in this repository with ``-DTUTORIAL_PYTHON=ON`` added to the CMake configuration options, -then install with ``cmake --build build --target pyamrex_pip_install``, and pyamrex will be installed for you. +then install with ``cmake --build build --target pyamrex_pip_install``, and pyAMReX will be installed for you. + +Running +======= + +Python tutorials are written so they run the same way as their C++ counterparts: +command line arguments after the script name are forwarded to AMReX, so an +:ref:`inputs file ` and ``key=value`` overrides can be passed as usual: + +.. code-block:: sh + + python3 main.py inputs + + # with runtime parameter override(s) + python3 main.py inputs nsteps=20 + + # MPI-parallel + mpiexec -n 2 python3 main.py inputs + +Tutorials +========= + +Guided tutorials: -Once pyAMReX is installed, you can run the following Guided Tutorial Examples: +- :download:`MultiFab <../../GuidedTutorials/MultiFab/main.py>`: define, fill and plot a MultiFab +- :download:`Heat Equation <../../GuidedTutorials/HeatEquation/Source/main.py>`: explicit heat equation solve with ghost cell exchanges and runtime parameters + (run with :download:`inputs <../../GuidedTutorials/HeatEquation/Exec/inputs>`) -- :download:`MultiFab <../../GuidedTutorials/MultiFab/main.py>` -- :download:`Heat Equation <../../GuidedTutorials/HeatEquation/Source/main.py>` +Example codes: +- :download:`MPMD Case-2 <../../ExampleCodes/MPMD/Case-2/main.py>`: a Python app coupled to a C++ app via AMReX MPMD (see :ref:`tutorials_mpmd`) diff --git a/ExampleCodes/MPMD/Case-2/main.py b/ExampleCodes/MPMD/Case-2/main.py index 542761f6..55f7adcd 100644 --- a/ExampleCodes/MPMD/Case-2/main.py +++ b/ExampleCodes/MPMD/Case-2/main.py @@ -1,3 +1,5 @@ +import sys + from mpi4py import MPI import amrex.space3d as amr @@ -5,11 +7,11 @@ # Initialize amrex::MPMD to establish communication across the two apps # However, leverage MPMD_Initialize_without_split # so that communication split can be performed using mpi4py.MPI -amr.MPMD_Initialize_without_split([]) +amr.MPMD_Initialize_without_split(sys.argv[1:]) # Leverage MPI from mpi4py to perform communication split app_comm_py = MPI.COMM_WORLD.Split(amr.MPMD_AppNum(), amr.MPMD_MyProc()) # Initialize AMReX -amr.initialize_when_MPMD([], app_comm_py) +amr.initialize_when_MPMD(sys.argv[1:], app_comm_py) amr.Print(f"Hello world from pyAMReX version {amr.__version__}\n") # Create a MPMD Copier that gets the BoxArray information from the other (C++) app diff --git a/GuidedTutorials/HeatEquation/Source/main.py b/GuidedTutorials/HeatEquation/Source/main.py index bc51812b..35cf74c2 100755 --- a/GuidedTutorials/HeatEquation/Source/main.py +++ b/GuidedTutorials/HeatEquation/Source/main.py @@ -8,6 +8,8 @@ # License: BSD-3-Clause-LBNL # Authors: Revathi Jambunathan, Edoardo Zoni, Olga Shapoval, David Grote, Axel Huebl +import sys + import amrex.space3d as amr @@ -96,8 +98,8 @@ def main(n_cell, max_grid_size, nsteps, plot_int, dt): # Loop over boxes for mfi in phi_old: bx = mfi.validbox() - # phiOld is indexed in reversed order (z,y,x) and indices are local - phiOld = xp.array(phi_old.array(mfi), copy=False) + # phiOld is indexed in reversed order (n,z,y,x) and indices are local + phiOld = phi_old.array(mfi).to_xp(copy=False, order="C") # set phi = 1 + e^(-(r-0.5)^2) x = (xp.arange(bx.small_end[0],bx.big_end[0]+1,1) + 0.5) * dx[0] y = (xp.arange(bx.small_end[1],bx.big_end[1]+1,1) + 0.5) * dx[1] @@ -121,8 +123,8 @@ def main(n_cell, max_grid_size, nsteps, plot_int, dt): # new_phi = old_phi + dt * Laplacian(old_phi) # Loop over boxes for mfi in phi_old: - phiOld = xp.array(phi_old.array(mfi), copy=False) - phiNew = xp.array(phi_new.array(mfi), copy=False) + phiOld = phi_old.array(mfi).to_xp(copy=False, order="C") + phiNew = phi_new.array(mfi).to_xp(copy=False, order="C") hix = phiOld.shape[3] hiy = phiOld.shape[2] hiz = phiOld.shape[1] @@ -157,20 +159,38 @@ def main(n_cell, max_grid_size, nsteps, plot_int, dt): if __name__ == '__main__': # Initialize AMReX - amr.initialize([]) - - # TODO Implement parser - # Simulation parameters - # number of cells on each side of the domain - n_cell = 32 - # size of each box (or grid) - max_grid_size = 16 - # total steps in simulation - nsteps = 1000 - # how often to write a plotfile - plot_int = 100 + # Command line arguments are forwarded, e.g., an inputs file: + # python3 main.py ../Exec/inputs + amr.initialize(sys.argv[1:]) + + # ********************************** + # SIMULATION PARAMETERS + + # ParmParse is a way of reading inputs from the inputs file + # pp.get means we require the inputs file to have it + # pp.query means we optionally need the inputs file to have it - but we must supply a default here + pp = amr.ParmParse() + + # We need to get n_cell from the inputs file - this is the number of cells on each side of + # a square (or cubic) domain. + n_cell = pp.get_int("n_cell") + + # The domain is broken into boxes of size max_grid_size + max_grid_size = pp.get_int("max_grid_size") + + # Default nsteps to 10, allow us to set it to something else in the inputs file + exists, nsteps = pp.query_int("nsteps") + if not exists: + nsteps = 10 + + # Default plot_int to -1, allow us to set it to something else in the inputs file + # If plot_int < 0 then no plot files will be written + exists, plot_int = pp.query_int("plot_int") + if not exists: + plot_int = -1 + # time step - dt = 1e-5 + dt = pp.get_real("dt") main(n_cell, max_grid_size, nsteps, plot_int, dt) diff --git a/GuidedTutorials/MultiFab/main.py b/GuidedTutorials/MultiFab/main.py index db1a2354..db1618e5 100755 --- a/GuidedTutorials/MultiFab/main.py +++ b/GuidedTutorials/MultiFab/main.py @@ -8,6 +8,8 @@ # License: BSD-3-Clause-LBNL # Authors: Revathi Jambunathan, Edoardo Zoni, Olga Shapoval, David Grote, Axel Huebl +import sys + import amrex.space3d as amr @@ -34,7 +36,8 @@ def load_cupy(): # Initialize AMReX -amr.initialize([]) +# Command line arguments after the script name are forwarded to AMReX +amr.initialize(sys.argv[1:]) # CPU/GPU logic xp = load_cupy() @@ -91,21 +94,23 @@ def load_cupy(): for mfi in mf: bx = mfi.validbox() # Preferred way to fill array using fast ranged operations: - # - xp.array is indexed in reversed order (n,z,y,x), - # .T creates a view into the AMReX (x,y,z,n) order + # - .to_xp() provides a NumPy (CPU) or CuPy (GPU) view into the + # data, indexed in the AMReX (x,y,z,n) order ("F" order, default) # - indices are local (range from 0 to box size) - mf_array = xp.array(mf.array(mfi), copy=False).T + mf_array = 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] - v = (x[xp.newaxis,xp.newaxis,:] - + y[xp.newaxis,:,xp.newaxis]*0.1 - + z[:,xp.newaxis,xp.newaxis]*0.01) rsquared = ((z[xp.newaxis, xp.newaxis, :] - 0.5)**2 + (y[xp.newaxis, :, xp.newaxis] - 0.5)**2 + (x[ :, xp.newaxis, xp.newaxis] - 0.5)**2) / 0.01 mf_array[:, :, :, 0] = 1. + xp.exp(-rsquared) +# Inspect the data: reductions over all boxes (and MPI ranks) +amr.Print(f"min(phi) = {mf.min(comp=0)}") +amr.Print(f"max(phi) = {mf.max(comp=0)}") +amr.Print(f"sum(phi) = {mf.sum(comp=0)}") + # Plot MultiFab data plotfile = amr.concatenate(root="plt", num=1, mindigits=3) varnames = amr.Vector_string(["comp0"])