From 6cf9dde3522bb1131895bf87c02cbfb52517dc2c Mon Sep 17 00:00:00 2001
From: Axel Huebl <axel.huebl@plasma.ninja>
Date: Sun, 5 Nov 2023 02:35:03 -0800
Subject: [PATCH] Pandas: `ImpactXParticleContainer.to_df()`

Copy all particles into a `pandas.DataFrame`.
Supports local and MPI-gathered results.
---
 src/python/ImpactXParticleContainer.cpp       | 48 ++++++++++++++++
 .../impactx/ImpactXParticleContainer.py       | 57 +++++++++++++++++++
 src/python/impactx/__init__.py                |  6 ++
 3 files changed, 111 insertions(+)
 create mode 100644 src/python/impactx/ImpactXParticleContainer.py

diff --git a/src/python/ImpactXParticleContainer.cpp b/src/python/ImpactXParticleContainer.cpp
index c30068d09..91561faed 100644
--- a/src/python/ImpactXParticleContainer.cpp
+++ b/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"),
diff --git a/src/python/impactx/ImpactXParticleContainer.py b/src/python/impactx/ImpactXParticleContainer.py
new file mode 100644
index 000000000..9fcb56f9d
--- /dev/null
+++ b/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
diff --git a/src/python/impactx/__init__.py b/src/python/impactx/__init__.py
index 4d4281b1c..15d2d799d 100644
--- a/src/python/impactx/__init__.py
+++ b/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)