Skip to content

Commit 4d3cb97

Browse files
ax3lclaude
andcommitted
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) <noreply@anthropic.com>
1 parent 624b081 commit 4d3cb97

5 files changed

Lines changed: 97 additions & 42 deletions

File tree

.github/workflows/linux.yml

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,10 @@ jobs:
3838
- name: Run Python
3939
run: |
4040
cd GuidedTutorials/MultiFab/
41-
python main.py
41+
python3 main.py
4242
cd ../HeatEquation/Source/
43-
python main.py
43+
python3 main.py ../Exec/inputs
44+
mpiexec -n 2 python3 main.py ../Exec/inputs
4445
4546
# Build all tutorials
4647
tutorials_omp:
@@ -69,9 +70,9 @@ jobs:
6970
- name: Run Python
7071
run: |
7172
cd GuidedTutorials/MultiFab/
72-
python main.py
73+
python3 main.py
7374
cd ../HeatEquation/Source/
74-
python main.py
75+
python3 main.py ../Exec/inputs
7576
7677
tutorials_clang:
7778
name: Clang SP Particles DP Mesh Debug [tutorials]
@@ -104,9 +105,9 @@ jobs:
104105
- name: Run Python
105106
run: |
106107
cd GuidedTutorials/MultiFab/
107-
python main.py
108+
python3 main.py
108109
cd ../HeatEquation/Source/
109-
python main.py
110+
python3 main.py ../Exec/inputs
110111
111112
# Build all tutorials w/o MPI
112113
tutorials-nonmpi:
@@ -135,9 +136,9 @@ jobs:
135136
- name: Run Python
136137
run: |
137138
cd GuidedTutorials/MultiFab/
138-
python main.py
139+
python3 main.py
139140
cd ../HeatEquation/Source/
140-
python main.py
141+
python3 main.py ../Exec/inputs
141142
142143
# Build all tutorials
143144
tutorials-nofortran:
@@ -164,9 +165,9 @@ jobs:
164165
- name: Run Python
165166
run: |
166167
cd GuidedTutorials/MultiFab/
167-
python main.py
168+
python3 main.py
168169
cd ../HeatEquation/Source/
169-
python main.py
170+
python3 main.py ../Exec/inputs
170171
171172
# Build all tutorials with CUDA
172173
tutorials-cuda:

Docs/source/Python_Tutorial.rst

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,44 @@
44
Python
55
======
66

7-
These examples show how to use AMReX from Python.
7+
These examples show how to use AMReX from Python, via `pyAMReX <https://github.com/AMReX-Codes/pyamrex/>`__.
88
AMReX applications can also be interfaced to Python with the same logic.
99

10+
Installation
11+
============
12+
1013
In order to run the Python tutorials, you need to have pyAMReX installed.
11-
Please see `pyAMReX <https://github.com/AMReX-Codes/pyamrex/>`__ for more details.
14+
Please see the `pyAMReX documentation <https://pyamrex.readthedocs.io>`__ for installation details.
1215

1316
Alternatively, you can build the ExampleCodes in this repository with ``-DTUTORIAL_PYTHON=ON`` added to the CMake configuration options,
14-
then install with ``cmake --build build --target pyamrex_pip_install``, and pyamrex will be installed for you.
17+
then install with ``cmake --build build --target pyamrex_pip_install``, and pyAMReX will be installed for you.
18+
19+
Running
20+
=======
21+
22+
Python tutorials are written so they run the same way as their C++ counterparts:
23+
command line arguments after the script name are forwarded to AMReX, so an
24+
:ref:`inputs file <amrex_docs:sec:basics:parmparse>` and ``key=value`` overrides can be passed as usual:
25+
26+
.. code-block:: sh
27+
28+
python3 main.py inputs
29+
30+
# with runtime parameter override(s)
31+
python3 main.py inputs nsteps=20
32+
33+
# MPI-parallel
34+
mpiexec -n 2 python3 main.py inputs
35+
36+
Tutorials
37+
=========
38+
39+
Guided tutorials:
1540

16-
Once pyAMReX is installed, you can run the following Guided Tutorial Examples:
41+
- :download:`MultiFab <../../GuidedTutorials/MultiFab/main.py>`: define, fill and plot a MultiFab
42+
- :download:`Heat Equation <../../GuidedTutorials/HeatEquation/Source/main.py>`: explicit heat equation solve with ghost cell exchanges and runtime parameters
43+
(run with :download:`inputs <../../GuidedTutorials/HeatEquation/Exec/inputs>`)
1744

18-
- :download:`MultiFab <../../GuidedTutorials/MultiFab/main.py>`
19-
- :download:`Heat Equation <../../GuidedTutorials/HeatEquation/Source/main.py>`
45+
Example codes:
2046

47+
- :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`)

ExampleCodes/MPMD/Case-2/main.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,17 @@
1+
import sys
2+
13
from mpi4py import MPI
24

35
import amrex.space3d as amr
46

57
# Initialize amrex::MPMD to establish communication across the two apps
68
# However, leverage MPMD_Initialize_without_split
79
# so that communication split can be performed using mpi4py.MPI
8-
amr.MPMD_Initialize_without_split([])
10+
amr.MPMD_Initialize_without_split(sys.argv[1:])
911
# Leverage MPI from mpi4py to perform communication split
1012
app_comm_py = MPI.COMM_WORLD.Split(amr.MPMD_AppNum(), amr.MPMD_MyProc())
1113
# Initialize AMReX
12-
amr.initialize_when_MPMD([], app_comm_py)
14+
amr.initialize_when_MPMD(sys.argv[1:], app_comm_py)
1315

1416
amr.Print(f"Hello world from pyAMReX version {amr.__version__}\n")
1517
# Create a MPMD Copier that gets the BoxArray information from the other (C++) app

GuidedTutorials/HeatEquation/Source/main.py

Lines changed: 37 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
# License: BSD-3-Clause-LBNL
99
# Authors: Revathi Jambunathan, Edoardo Zoni, Olga Shapoval, David Grote, Axel Huebl
1010

11+
import sys
12+
1113
import amrex.space3d as amr
1214

1315

@@ -96,8 +98,8 @@ def main(n_cell, max_grid_size, nsteps, plot_int, dt):
9698
# Loop over boxes
9799
for mfi in phi_old:
98100
bx = mfi.validbox()
99-
# phiOld is indexed in reversed order (z,y,x) and indices are local
100-
phiOld = xp.array(phi_old.array(mfi), copy=False)
101+
# phiOld is indexed in reversed order (n,z,y,x) and indices are local
102+
phiOld = phi_old.array(mfi).to_xp(copy=False, order="C")
101103
# set phi = 1 + e^(-(r-0.5)^2)
102104
x = (xp.arange(bx.small_end[0],bx.big_end[0]+1,1) + 0.5) * dx[0]
103105
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):
121123
# new_phi = old_phi + dt * Laplacian(old_phi)
122124
# Loop over boxes
123125
for mfi in phi_old:
124-
phiOld = xp.array(phi_old.array(mfi), copy=False)
125-
phiNew = xp.array(phi_new.array(mfi), copy=False)
126+
phiOld = phi_old.array(mfi).to_xp(copy=False, order="C")
127+
phiNew = phi_new.array(mfi).to_xp(copy=False, order="C")
126128
hix = phiOld.shape[3]
127129
hiy = phiOld.shape[2]
128130
hiz = phiOld.shape[1]
@@ -157,20 +159,38 @@ def main(n_cell, max_grid_size, nsteps, plot_int, dt):
157159

158160
if __name__ == '__main__':
159161
# Initialize AMReX
160-
amr.initialize([])
161-
162-
# TODO Implement parser
163-
# Simulation parameters
164-
# number of cells on each side of the domain
165-
n_cell = 32
166-
# size of each box (or grid)
167-
max_grid_size = 16
168-
# total steps in simulation
169-
nsteps = 1000
170-
# how often to write a plotfile
171-
plot_int = 100
162+
# Command line arguments are forwarded, e.g., an inputs file:
163+
# python3 main.py ../Exec/inputs
164+
amr.initialize(sys.argv[1:])
165+
166+
# **********************************
167+
# SIMULATION PARAMETERS
168+
169+
# ParmParse is a way of reading inputs from the inputs file
170+
# pp.get means we require the inputs file to have it
171+
# pp.query means we optionally need the inputs file to have it - but we must supply a default here
172+
pp = amr.ParmParse()
173+
174+
# We need to get n_cell from the inputs file - this is the number of cells on each side of
175+
# a square (or cubic) domain.
176+
n_cell = pp.get_int("n_cell")
177+
178+
# The domain is broken into boxes of size max_grid_size
179+
max_grid_size = pp.get_int("max_grid_size")
180+
181+
# Default nsteps to 10, allow us to set it to something else in the inputs file
182+
exists, nsteps = pp.query_int("nsteps")
183+
if not exists:
184+
nsteps = 10
185+
186+
# Default plot_int to -1, allow us to set it to something else in the inputs file
187+
# If plot_int < 0 then no plot files will be written
188+
exists, plot_int = pp.query_int("plot_int")
189+
if not exists:
190+
plot_int = -1
191+
172192
# time step
173-
dt = 1e-5
193+
dt = pp.get_real("dt")
174194

175195
main(n_cell, max_grid_size, nsteps, plot_int, dt)
176196

GuidedTutorials/MultiFab/main.py

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
# License: BSD-3-Clause-LBNL
99
# Authors: Revathi Jambunathan, Edoardo Zoni, Olga Shapoval, David Grote, Axel Huebl
1010

11+
import sys
12+
1113
import amrex.space3d as amr
1214

1315

@@ -34,7 +36,8 @@ def load_cupy():
3436

3537

3638
# Initialize AMReX
37-
amr.initialize([])
39+
# Command line arguments after the script name are forwarded to AMReX
40+
amr.initialize(sys.argv[1:])
3841

3942
# CPU/GPU logic
4043
xp = load_cupy()
@@ -91,21 +94,23 @@ def load_cupy():
9194
for mfi in mf:
9295
bx = mfi.validbox()
9396
# Preferred way to fill array using fast ranged operations:
94-
# - xp.array is indexed in reversed order (n,z,y,x),
95-
# .T creates a view into the AMReX (x,y,z,n) order
97+
# - .to_xp() provides a NumPy (CPU) or CuPy (GPU) view into the
98+
# data, indexed in the AMReX (x,y,z,n) order ("F" order, default)
9699
# - indices are local (range from 0 to box size)
97-
mf_array = xp.array(mf.array(mfi), copy=False).T
100+
mf_array = mf.array(mfi).to_xp(copy=False)
98101
x = (xp.arange(bx.small_end[0], bx.big_end[0]+1)+0.5)*dx[0]
99102
y = (xp.arange(bx.small_end[1], bx.big_end[1]+1)+0.5)*dx[1]
100103
z = (xp.arange(bx.small_end[2], bx.big_end[2]+1)+0.5)*dx[2]
101-
v = (x[xp.newaxis,xp.newaxis,:]
102-
+ y[xp.newaxis,:,xp.newaxis]*0.1
103-
+ z[:,xp.newaxis,xp.newaxis]*0.01)
104104
rsquared = ((z[xp.newaxis, xp.newaxis, :] - 0.5)**2
105105
+ (y[xp.newaxis, :, xp.newaxis] - 0.5)**2
106106
+ (x[ :, xp.newaxis, xp.newaxis] - 0.5)**2) / 0.01
107107
mf_array[:, :, :, 0] = 1. + xp.exp(-rsquared)
108108

109+
# Inspect the data: reductions over all boxes (and MPI ranks)
110+
amr.Print(f"min(phi) = {mf.min(comp=0)}")
111+
amr.Print(f"max(phi) = {mf.max(comp=0)}")
112+
amr.Print(f"sum(phi) = {mf.sum(comp=0)}")
113+
109114
# Plot MultiFab data
110115
plotfile = amr.concatenate(root="plt", num=1, mindigits=3)
111116
varnames = amr.Vector_string(["comp0"])

0 commit comments

Comments
 (0)