Doc: Simple Compute (Helpers) (#273)

ax3l · web-flow · commit a44b53af40e0 · 2024-03-05T19:12:04.000-08:00
Highlight in the manual how to use NumPy/CuPy/Pandas helpers on whole
data containers for simple compute operations.
diff --git a/docs/source/usage/compute.rst b/docs/source/usage/compute.rst
@@ -30,11 +30,24 @@ The field data can have more than one component (in the slowest varying index),
 
 This is how to iterate and potentially compute for all blocks assigned to a local process in pyAMReX:
 
-.. literalinclude:: ../../../tests/test_multifab.py
-   :language: python3
-   :dedent: 4
-   :start-after: # Manual: Compute Mfab START
-   :end-before: # Manual: Compute Mfab END
+.. tab-set::
+
+   .. tab-item:: Simple
+
+      .. literalinclude:: ../../../tests/test_multifab.py
+         :language: python3
+         :dedent: 4
+         :start-after: # Manual: Compute Mfab Simple START
+         :end-before: # Manual: Compute Mfab Simple END
+
+   .. tab-item:: Detailed
+
+      .. literalinclude:: ../../../tests/test_multifab.py
+         :language: python3
+         :dedent: 4
+         :start-after: # Manual: Compute Mfab Detailed START
+         :end-before: # Manual: Compute Mfab Detailed END
+
 
 For a complete physics example that uses CPU/GPU agnostic Python code for computation on fields, see:
 
@@ -69,19 +82,31 @@ Here is the general structure for computing on particles:
 
    .. tab-item:: Modern (pure SoA) Layout
 
-      .. literalinclude:: ../../../tests/test_particleContainer.py
-         :language: python3
-         :dedent: 4
-         :start-after: # Manual: Pure SoA Compute PC START
-         :end-before: # Manual: Pure SoA Compute PC END
+      .. tab-set::
+
+         .. tab-item:: Simple: Pandas (read-only)
+
+            .. literalinclude:: ../../../tests/test_particleContainer.py
+               :language: python3
+               :dedent: 4
+               :start-after: # Manual: Pure SoA Compute PC Pandas START
+               :end-before: # Manual: Pure SoA Compute PC Pandas END
+
+         .. tab-item:: Detailed (read and write)
+
+            .. literalinclude:: ../../../tests/test_particleContainer.py
+               :language: python3
+               :dedent: 4
+               :start-after: # Manual: Pure SoA Compute PC Detailed START
+               :end-before: # Manual: Pure SoA Compute PC Detailed END
 
    .. tab-item:: Legacy (AoS + SoA) Layout
 
       .. literalinclude:: ../../../tests/test_particleContainer.py
          :language: python3
          :dedent: 4
-         :start-after: # Manual: Legacy Compute PC START
-         :end-before: # Manual: Legacy Compute PC END
+         :start-after: # Manual: Legacy Compute PC Detailed START
+         :end-before: # Manual: Legacy Compute PC Detailed END
 
 For many small CPU and GPU examples on how to compute on particles, see the following test cases:
 
diff --git a/tests/test_multifab.py b/tests/test_multifab.py
@@ -11,39 +11,64 @@
 def test_mfab_numpy(mfab):
     """Used in docs/source/usage/compute.rst"""
 
-    # Manual: Compute Mfab START
-    # finest active MR level, get from a simulation's AmrMesh object, e.g.:
+    class Config:
+        have_gpu = False
+
+    # Manual: Compute Mfab Detailed START
+    # finest active MR level, get from a
+    # simulation's AmrMesh object, e.g.:
     # finest_level = sim.finest_level
     finest_level = 0  # no MR
 
-    # iterate over every mesh-refinement levels
+    # iterate over mesh-refinement levels
     for lev in range(finest_level + 1):
-        # get an existing MultiFab, e.g., from a simulation:
-        # mfab = sim.get_field(lev=lev)  # code-specific getter function
+        # get an existing MultiFab, e.g.,
+        # from a simulation:
+        # mfab = sim.get_field(lev=lev)
+        # Config = sim.extension.Config
 
         # grow (aka guard/ghost/halo) regions
         ngv = mfab.n_grow_vect
 
         # get every local block of the field
         for mfi in mfab:
-            # global index space box, including guards
+            # global index box w/ guards
             bx = mfi.tilebox().grow(ngv)
-            print(bx)  # note: global index space of this block
-
-            # numpy representation: non-copying view, including the
-            # guard/ghost region
-            field_np = mfab.array(mfi).to_numpy()
+            print(bx)
+
+            # numpy representation: non-
+            # copying view, w/ guard/ghost
+            field = (
+                mfab.array(mfi).to_cupy()
+                if Config.have_gpu
+                else mfab.array(mfi).to_numpy()
+            )
 
-            # notes on indexing in field_np:
+            # notes on indexing in field:
             # - numpy uses locally zero-based indexing
             # - layout is F_CONTIGUOUS by default, just like AMReX
 
-            # notes:
-            # Only the next lines are the "HOT LOOP" of the computation.
-            # For efficiency, use numpy array operation for speed on CPUs.
-            # For GPUs use .to_cupy() above and compute with cupy or numba.
-            field_np[()] = 42.0
-    # Manual: Compute Mfab END
+            field[()] = 42.0
+    # Manual: Compute Mfab Detailed END
+
+    # Manual: Compute Mfab Simple START
+    # finest active MR level, get from a
+    # simulation's AmrMesh object, e.g.:
+    # finest_level = sim.finest_level
+    finest_level = 0  # no MR
+
+    # iterate over mesh-refinement levels
+    for lev in range(finest_level + 1):
+        # get an existing MultiFab, e.g.,
+        # from a simulation:
+        # mfab = sim.get_field(lev=lev)
+        # Config = sim.extension.Config
+
+        field_list = mfab.to_cupy() if Config.have_gpu else mfab.to_numpy()
+
+        for field in field_list:
+            field[()] = 42.0
+    # Manual: Compute Mfab Simple END
 
 
 def test_mfab_loop(mfab):
diff --git a/tests/test_particleContainer.py b/tests/test_particleContainer.py
@@ -319,24 +319,23 @@ def test_soa_pc_numpy(soa_particle_container, Npart):
     class Config:
         have_gpu = False
 
-    # Manual: Pure SoA Compute PC START
+    # Manual: Pure SoA Compute PC Detailed START
     # code-specific getter function, e.g.:
     # pc = sim.get_particles()
     # Config = sim.extension.Config
 
-    # iterate over mesh-refinement level (no MR: lev=0)
+    # iterate over mesh-refinement levels
     for lvl in range(pc.finest_level + 1):
-        # get every local chunk of particles
+        # loop local tiles of particles
         for pti in pc.iterator(pc, level=lvl):
-            # compile-time and runtime attributes in SoA format
+            # compile-time and runtime attributes
             soa = pti.soa().to_cupy() if Config.have_gpu else pti.soa().to_numpy()
 
             # print all particle ids in the chunk
             print("idcpu =", soa.idcpu)
 
-            # notes:
-            # Only the next lines are the "HOT LOOP" of the computation.
-            # For speed, use array operations.
+            x = soa.real["x"]
+            y = soa.real["y"]
 
             # write to all particles in the chunk
             # note: careful, if you change particle positions, you might need to
@@ -345,15 +344,15 @@ class Config:
             soa.real["y"][:] = 0.35
             soa.real["z"][:] = 0.40
 
-            soa.real["a"][:] = 0.50
-            soa.real["b"][:] = 0.50
+            soa.real["a"][:] = x[:] ** 2
+            soa.real["b"][:] = x[:] + y[:]
             soa.real["c"][:] = 0.50
             # ...
 
             # all int attributes
             for soa_int in soa.int.values():
                 soa_int[:] = 12
-    # Manual: Pure SoA Compute PC END
+    # Manual: Pure SoA Compute PC Detailed END
 
 
 def test_pc_numpy(particle_container, Npart):
@@ -363,14 +362,14 @@ def test_pc_numpy(particle_container, Npart):
     class Config:
         have_gpu = False
 
-    # Manual: Legacy Compute PC START
+    # Manual: Legacy Compute PC Detailed START
     # code-specific getter function, e.g.:
     # pc = sim.get_particles()
     # Config = sim.extension.Config
 
-    # iterate over mesh-refinement level (no MR: lev=0)
+    # iterate over mesh-refinement levels
     for lvl in range(pc.finest_level + 1):
-        # get every local chunk of particles
+        # loop local tiles of particles
         for pti in pc.iterator(pc, level=lvl):
             # default layout: AoS with positions and idcpu
             # note: not part of the new PureSoA particle container layout
@@ -402,7 +401,7 @@ class Config:
 
             for soa_int in soa.int.values():
                 soa_int[:] = 12
-    # Manual: Legacy Compute PC END
+    # Manual: Legacy Compute PC Detailed END
 
 
 @pytest.mark.skipif(
@@ -440,6 +439,47 @@ def test_soa_pc_df_mpi(soa_particle_container, Npart):
         print(df)
 
 
+@pytest.mark.skipif(
+    importlib.util.find_spec("pandas") is None, reason="pandas is not available"
+)
+def test_soa_pc_df(soa_particle_container, Npart):
+    """Used in docs/source/usage/compute.rst"""
+    pc = soa_particle_container
+
+    class Config:
+        have_gpu = False
+
+    # Manual: Pure SoA Compute PC Pandas START
+    # code-specific getter function, e.g.:
+    # pc = sim.get_particles()
+    # Config = sim.extension.Config
+
+    # local particles on all levels
+    df = pc.to_df()  # this is a copy!
+    print(df)
+
+    # read
+    print(df["x"])
+
+    # write (into copy!)
+    df["x"] = 0.30
+    df["y"] = 0.35
+    df["z"] = 0.40
+
+    df["a"] = df["x"] ** 2
+    df["b"] = df["x"] + df["y"]
+    df["c"] = 0.50
+
+    # int attributes
+    # df["i1"] = 12
+    # df["i2"] = 12
+    # ...
+
+    print(df)
+
+    # Manual: Pure SoA Compute PC Pandas END
+
+
 @pytest.mark.skipif(
     importlib.util.find_spec("pandas") is None, reason="pandas is not available"
 )
