Pandas: ParticleContainer_*.to_df()

Copy all particles into a `pandas.DataFrame`. Supports local and
MPI-gathered results.

Author: ax3l

src/Particle/ParticleContainer.H

@@ -68,6 +68,8 @@ void make_Base_Iterators (py::module &m, std::string allocstr)
                                py::return_value_policy::reference_internal)
 
         .def_property_readonly_static("is_soa_particle", [](const py::object&){ return ParticleType::is_soa_particle;})
+        .def_property_readonly("size", &iterator_base::numParticles,
+                               "the number of particles on this tile")
         .def_property_readonly("num_particles", &iterator_base::numParticles)
         .def_property_readonly("num_real_particles", &iterator_base::numRealParticles)
         .def_property_readonly("num_neighbor_particles", &iterator_base::numNeighborParticles)
@@ -382,6 +384,14 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
     make_Iterators< false, iterator, Allocator >(m, allocstr);
     using const_iterator = amrex::ParConstIter_impl<ParticleType, T_NArrayReal, T_NArrayInt, Allocator>;
     make_Iterators< true, const_iterator, Allocator >(m, allocstr);
+
+    // simpler particle iterator loops: return types of this particle box
+    py_pc
+        .def_property_readonly_static("iterator", [](py::object /* pc */){ return py::type::of<iterator>(); },
+            "amrex iterator for particle boxes")
+        .def_property_readonly_static("const_iterator", [](py::object /* pc */){ return py::type::of<const_iterator>(); },
+            "amrex constant iterator for particle boxes (read-only)")
+    ;
 }
 
 /** Create ParticleContainers and Iterators

src/amrex/ParticleContainer.py

@@ -0,0 +1,104 @@
+"""
+This file is part of pyAMReX
+
+Copyright 2023 AMReX community
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def pc_to_df(self, local=True, comm=None, root_rank=0):
+    """
+    Copy all particles into a pandas.DataFrame
+
+    Parameters
+    ----------
+    self : amrex.ParticleContainer_*
+        A ParticleContainer class in pyAMReX
+    local : bool
+        MPI-local particles
+    comm : MPI Communicator
+        if local is False, this defaults to mpi4py.MPI.COMM_WORLD
+    root_rank : MPI root rank to gather to
+        if local is False, this defaults to 0
+
+    Returns
+    -------
+    A concatenated pandas.DataFrame with particles from all levels.
+
+    Returns None if no particles were found.
+    If local=False, then all ranks but the root_rank will return None.
+    """
+    import pandas as pd
+
+    # create a DataFrame per particle box and append it to the list of
+    # local DataFrame(s)
+    dfs_local = []
+    for lvl in range(self.finest_level + 1):
+        for pti in self.const_iterator(self, level=lvl):
+            if pti.size == 0:
+                continue
+
+            if self.is_soa_particle:
+                next_df = pd.DataFrame()
+            else:
+                # AoS
+                aos_np = pti.aos().to_numpy(copy=True)
+                next_df = pd.DataFrame(aos_np)
+                next_df.set_index("cpuid")
+                next_df.index.name = "cpuid"
+
+            # SoA
+            soa_view = pti.soa().to_numpy(copy=True)
+            soa_np_real = soa_view.real
+            soa_np_int = soa_view.int
+
+            for idx, array in enumerate(soa_np_real):
+                next_df[f"SoA_real_{idx}"] = array
+            for idx, array in enumerate(soa_np_int):
+                next_df[f"SoA_int_{idx}"] = array
+
+            dfs_local.append(next_df)
+
+    # MPI Gather to root rank if requested
+    if local:
+        if len(dfs_local) == 0:
+            df = None
+        else:
+            df = pd.concat(dfs_local)
+    else:
+        from mpi4py import MPI
+
+        if comm is None:
+            comm = MPI.COMM_WORLD
+        rank = comm.Get_rank()
+
+        # a list for each rank's list of DataFrame(s)
+        df_list_list = comm.gather(dfs_local, root=root_rank)
+
+        if rank == root_rank:
+            flattened_list = [df for sublist in df_list_list for df in sublist]
+
+            if len(flattened_list) == 0:
+                df = pd.DataFrame()
+            else:
+                df = pd.concat(flattened_list, ignore_index=True)
+        else:
+            df = None
+
+    return df
+
+
+def register_ParticleContainer_extension(amr):
+    """ParticleContainer helper methods"""
+    import inspect
+    import sys
+
+    # register member functions for every ParticleContainer_* type
+    for _, ParticleContainer_type in inspect.getmembers(
+        sys.modules[amr.__name__],
+        lambda member: inspect.isclass(member)
+        and member.__module__ == amr.__name__
+        and member.__name__.startswith("ParticleContainer_"),
+    ):
+        ParticleContainer_type.to_df = pc_to_df

src/amrex/space1d/__init__.py

@@ -48,10 +48,12 @@ def Print(*args, **kwargs):
 from ..ArrayOfStructs import register_AoS_extension
 from ..MultiFab import register_MultiFab_extension
 from ..PODVector import register_PODVector_extension
+from ..ParticleContainer import register_ParticleContainer_extension
 from ..StructOfArrays import register_SoA_extension
 
 register_Array4_extension(amrex_1d_pybind)
 register_MultiFab_extension(amrex_1d_pybind)
 register_PODVector_extension(amrex_1d_pybind)
 register_SoA_extension(amrex_1d_pybind)
 register_AoS_extension(amrex_1d_pybind)
+register_ParticleContainer_extension(amrex_1d_pybind)

src/amrex/space2d/__init__.py

@@ -48,10 +48,12 @@ def Print(*args, **kwargs):
 from ..ArrayOfStructs import register_AoS_extension
 from ..MultiFab import register_MultiFab_extension
 from ..PODVector import register_PODVector_extension
+from ..ParticleContainer import register_ParticleContainer_extension
 from ..StructOfArrays import register_SoA_extension
 
 register_Array4_extension(amrex_2d_pybind)
 register_MultiFab_extension(amrex_2d_pybind)
 register_PODVector_extension(amrex_2d_pybind)
 register_SoA_extension(amrex_2d_pybind)
 register_AoS_extension(amrex_2d_pybind)
+register_ParticleContainer_extension(amrex_2d_pybind)

src/amrex/space3d/__init__.py

@@ -48,10 +48,12 @@ def Print(*args, **kwargs):
 from ..ArrayOfStructs import register_AoS_extension
 from ..MultiFab import register_MultiFab_extension
 from ..PODVector import register_PODVector_extension
+from ..ParticleContainer import register_ParticleContainer_extension
 from ..StructOfArrays import register_SoA_extension
 
 register_Array4_extension(amrex_3d_pybind)
 register_MultiFab_extension(amrex_3d_pybind)
 register_PODVector_extension(amrex_3d_pybind)
 register_SoA_extension(amrex_3d_pybind)
 register_AoS_extension(amrex_3d_pybind)
+register_ParticleContainer_extension(amrex_3d_pybind)

tests/test_particleContainer.py

@@ -272,3 +272,29 @@ def test_per_cell(empty_particle_container, std_geometry, std_particle):
     assert pc.TotalNumberOfParticles() == pc.NumberOfParticlesAtLevel(0) == ncells
     print("npts * real_1", ncells * std_particle.real_array_data[1])
     assert ncells * std_particle.real_array_data[1] == sum_1
+
+
+def test_pc_df(particle_container, Npart):
+    pc = particle_container
+    print(f"pc={pc}")
+    df = pc.to_df()
+    print(df.columns)
+    print(df)
+
+
+def test_pc_empty_df(empty_particle_container, Npart):
+    pc = empty_particle_container
+    print(f"pc={pc}")
+    df = pc.to_df()
+    assert df is None
+
+
+@pytest.mark.skipif(not amr.Config.have_mpi, reason="Requires AMReX_MPI=ON")
+def test_pc_df_mpi(particle_container, Npart):
+    pc = particle_container
+    print(f"pc={pc}")
+    df = pc.to_df(local=False)
+    if df is not None:
+        # only rank 0
+        print(df.columns)
+        print(df)