Pandas: ImpactXParticleContainer.to_df() (#458)

* Pandas: `ImpactXParticleContainer.to_df()`

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

* AMReX & pyAMReX: latest development

X-ref:
AMReX-Codes/pyamrex#220

* DataFrame: Add Test

Author: ax3l

cmake/dependencies/ABLASTR.cmake

@@ -186,7 +186,7 @@ set(ImpactX_ablastr_branch "23.11"
 set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git"
     CACHE STRING
     "Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)")
-set(ImpactX_amrex_branch ""
+set(ImpactX_amrex_branch "175b99d913dc2748e43c53192737170c770fe0e8"
     CACHE STRING
     "Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)")
 

cmake/dependencies/pyAMReX.cmake

@@ -79,7 +79,7 @@ option(ImpactX_pyamrex_internal "Download & build pyAMReX" ON)
 set(ImpactX_pyamrex_repo "https://github.com/AMReX-Codes/pyamrex.git"
     CACHE STRING
     "Repository URI to pull and build pyamrex from if(ImpactX_pyamrex_internal)")
-set(ImpactX_pyamrex_branch "23.11"
+set(ImpactX_pyamrex_branch "1c24a8504fd9df0e2ad2f447a4d1567750a0e4e0"
     CACHE STRING
     "Repository branch for ImpactX_pyamrex_repo if(ImpactX_pyamrex_internal)")
 

src/python/ImpactXParticleContainer.cpp

@@ -12,12 +12,50 @@
 #include <AMReX_MFIter.H>
 #include <AMReX_ParticleContainer.H>
 
+#include <algorithm>
+#include <string>
+#include <vector>
+
 namespace py = pybind11;
 using namespace impactx;
 
 
 void init_impactxparticlecontainer(py::module& m)
 {
+    py::class_<RealAoS>(m, "RealAoS")
+        .def_property_readonly_static("names_s",
+            [](py::object) {
+                std::vector<std::string> real_aos_names(RealAoS::names_s.size());
+                std::copy(RealAoS::names_s.begin(), RealAoS::names_s.end(), real_aos_names.begin());
+                return real_aos_names;
+            },
+            "named labels for fixed s")
+        .def_property_readonly_static("names_t",
+            [](py::object) {
+                std::vector<std::string> real_aos_names(RealAoS::names_t.size());
+                std::copy(RealAoS::names_t.begin(), RealAoS::names_t.end(), real_aos_names.begin());
+                return real_aos_names;
+            },
+            "named labels for fixed t")
+    ;
+
+    py::class_<RealSoA>(m, "RealSoA")
+        .def_property_readonly_static("names_s",
+        [](py::object) {
+                std::vector<std::string> real_soa_names(RealSoA::names_s.size());
+                std::copy(RealSoA::names_s.begin(), RealSoA::names_s.end(), real_soa_names.begin());
+                return real_soa_names;
+        },
+        "named labels for fixed s")
+        .def_property_readonly_static("names_t",
+    [](py::object) {
+            std::vector<std::string> real_soa_names(RealSoA::names_t.size());
+            std::copy(RealSoA::names_t.begin(), RealSoA::names_t.end(), real_soa_names.begin());
+            return real_soa_names;
+        },
+        "named labels for fixed t")
+    ;
+
     py::class_<
         ParIter,
         amrex::ParIter<0, 0, RealSoA::nattribs, IntSoA::nattribs>
@@ -43,6 +81,16 @@ void init_impactxparticlecontainer(py::module& m)
         amrex::ParticleContainer<0, 0, RealSoA::nattribs, IntSoA::nattribs>
     >(m, "ImpactXParticleContainer")
         //.def(py::init<>())
+
+        .def_property_readonly_static("RealAoS",
+            [](py::object /* pc */){ return py::type::of<RealAoS>(); },
+            "RealAoS attribute name labels"
+        )
+        .def_property_readonly_static("RealSoA",
+            [](py::object /* pc */){ return py::type::of<RealSoA>(); },
+            "RealSoA attribute name labels"
+        )
+
         .def("add_n_particles",
              &ImpactXParticleContainer::AddNParticles,
              py::arg("lev"),

src/python/impactx/ImpactXParticleContainer.py

@@ -0,0 +1,57 @@
+"""
+This file is part of ImpactX
+
+Copyright 2023 ImpactX contributors
+Authors: Axel Huebl
+License: BSD-3-Clause-LBNL
+"""
+
+
+def ix_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.
+    """
+    df = super(type(self), self).to_df(local=local, comm=comm, root_rank=root_rank)
+
+    # rename columns according to our attribute names
+    if df is not None:
+        # todo: check if currently in fixed s or fixed t and pick name accordingly
+
+        names = []
+        for n in self.RealAoS.names_s:
+            names.append(n)
+        names.append("cpuid")
+        for n in self.RealSoA.names_s:
+            names.append(n)
+
+        df.columns.values[0 : len(names)] = names
+
+        # todo: also rename runtime attributes (e.g., "s_lost")
+        # https://github.com/ECP-WarpX/impactx/pull/398
+
+    return df
+
+
+def register_ImpactXParticleContainer_extension(ixpc):
+    """ImpactXParticleContainer helper methods"""
+
+    # register member functions for ImpactXParticleContainer
+    ixpc.to_df = ix_pc_to_df

src/python/impactx/__init__.py

@@ -17,6 +17,9 @@
 
 # import core bindings to C++
 from . import impactx_pybind as cxx
+from .ImpactXParticleContainer import (
+    register_ImpactXParticleContainer_extension,
+)
 from .impactx_pybind import *  # noqa
 from .madx_to_impactx import read_beam, read_lattice  # noqa
 
@@ -35,3 +38,6 @@
 
 # MAD-X file reader for reference particle
 RefPart.load_file = read_beam  # noqa
+
+# Pure Python extensions to ImpactX types
+register_ImpactXParticleContainer_extension(ImpactXParticleContainer)

tests/python/test_dataframe.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 The ImpactX Community
+#
+# Authors: Axel Huebl
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from impactx import ImpactX, RefPart, distribution, elements
+
+
+def test_df_pandas():
+    """
+    This tests using ImpactX and Pandas Dataframes
+    """
+    sim = ImpactX()
+
+    sim.particle_shape = 2
+    sim.slice_step_diagnostics = True
+    sim.init_grids()
+
+    # init particle beam
+    kin_energy_MeV = 2.0e3
+    bunch_charge_C = 1.0e-9
+    npart = 10000
+
+    #   reference particle
+    pc = sim.particle_container()
+    ref = pc.ref_particle()
+    ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+    #   particle bunch
+    distr = distribution.Waterbag(
+        sigmaX=3.9984884770e-5,
+        sigmaY=3.9984884770e-5,
+        sigmaT=1.0e-3,
+        sigmaPx=2.6623538760e-5,
+        sigmaPy=2.6623538760e-5,
+        sigmaPt=2.0e-3,
+        muxpx=-0.846574929020762,
+        muypy=0.846574929020762,
+        mutpt=0.0,
+    )
+    sim.add_particles(bunch_charge_C, distr, npart)
+
+    assert pc.TotalNumberOfParticles() == npart
+
+    # init accelerator lattice
+    fodo = [
+        elements.Drift(0.25),
+        elements.Quad(1.0, 1.0),
+        elements.Drift(0.5),
+        elements.Quad(1.0, -1.0),
+        elements.Drift(0.25),
+    ]
+    sim.lattice.extend(fodo)
+
+    # simulate
+    sim.evolve()
+
+    # compare number of global particles
+    df = pc.to_df(local=False)
+    if df is not None:
+        assert npart == len(df)
+
+
+if __name__ == "__main__":
+    test_df_pandas()