diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index 0c931fe3d..03b416a40 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/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)") diff --git a/cmake/dependencies/pyAMReX.cmake b/cmake/dependencies/pyAMReX.cmake index 3964fc373..98589d23c 100644 --- a/cmake/dependencies/pyAMReX.cmake +++ b/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)") 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 #include +#include +#include +#include + namespace py = pybind11; using namespace impactx; void init_impactxparticlecontainer(py::module& m) { + py::class_(m, "RealAoS") + .def_property_readonly_static("names_s", + [](py::object) { + std::vector 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 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_(m, "RealSoA") + .def_property_readonly_static("names_s", + [](py::object) { + std::vector 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 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 attribute name labels" + ) + .def_property_readonly_static("RealSoA", + [](py::object /* pc */){ return py::type::of(); }, + "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) diff --git a/tests/python/test_dataframe.py b/tests/python/test_dataframe.py new file mode 100644 index 000000000..a209317b4 --- /dev/null +++ b/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()