Python - add new particle attributes at runtime (#2112)

* exposes AddRealComp to Python to allow extra particle attributes to be added at runtime; also includes a new function to grab a particle data array from the name of the component rather than the index

* added function to get index of a particle component given the PID name

* changed new get component index and get_particle_arrays_from_comp_name functions to take species name as argument rather than species id

* changed warpx_addRealComp to accept a species name as input and only add the new component for that species

2 files changed, 65 insertions, 0 deletions

diff --git a/Source/Python/WarpXWrappers.cpp b/Source/Python/WarpXWrappers.cpp
index 3a7a91a55..618829fb5 100644
--- a/Source/Python/WarpXWrappers.cpp
+++ b/Source/Python/WarpXWrappers.cpp
@@ -404,6 +404,29 @@ extern "C"
         const auto & mypc = WarpX::GetInstance().GetPartContainer();
         auto & myspc = mypc.GetParticleContainer(speciesnumber);
 
+        return warpx_getParticleArraysUsingPC(
+            myspc, comp, lev, num_tiles, particles_per_tile
+        );
+    }
+
+    amrex::ParticleReal** warpx_getParticleArraysFromCompName (
+            const char* char_species_name, const char* char_comp_name,
+            int lev, int* num_tiles, int** particles_per_tile ) {
+
+        const auto & mypc = WarpX::GetInstance().GetPartContainer();
+        const std::string species_name(char_species_name);
+        auto & myspc = mypc.GetParticleContainerFromName(species_name);
+
+        return warpx_getParticleArraysUsingPC(
+            myspc,
+            warpx_getParticleCompIndex(char_species_name, char_comp_name), lev,
+            num_tiles, particles_per_tile
+        );
+    }
+
+    amrex::ParticleReal** warpx_getParticleArraysUsingPC (
+            WarpXParticleContainer& myspc, int comp,
+            int lev, int* num_tiles, int** particles_per_tile ) {
         int i = 0;
         for (WarpXParIter pti(myspc, lev); pti.isValid(); ++pti, ++i) {}
 
@@ -421,6 +444,33 @@ extern "C"
         return data;
     }
 
+    int warpx_getParticleCompIndex (
+         const char* char_species_name, const char* char_comp_name )
+    {
+        const auto & mypc = WarpX::GetInstance().GetPartContainer();
+
+        const std::string species_name(char_species_name);
+        auto & myspc = mypc.GetParticleContainerFromName(species_name);
+
+        const std::string comp_name(char_comp_name);
+        auto particle_comps = myspc.getParticleComps();
+
+        return particle_comps.at(comp_name);
+    }
+
+    void warpx_addRealComp(const char* char_species_name,
+        const char* char_comp_name, bool comm=true)
+    {
+        auto & mypc = WarpX::GetInstance().GetPartContainer();
+        const std::string species_name(char_species_name);
+        auto & myspc = mypc.GetParticleContainerFromName(species_name);
+
+        const std::string comp_name(char_comp_name);
+        myspc.AddRealComp(comp_name, comm);
+
+        mypc.defineAllParticleTiles();
+    }
+
     void warpx_ComputeDt () {
         WarpX& warpx = WarpX::GetInstance();
         warpx.ComputeDt ();
diff --git a/Source/Python/WarpXWrappers.h b/Source/Python/WarpXWrappers.h
index f98600707..1672b01ae 100644
--- a/Source/Python/WarpXWrappers.h
+++ b/Source/Python/WarpXWrappers.h
@@ -8,6 +8,7 @@
 #ifndef WARPX_WRAPPERS_H_
 #define WARPX_WRAPPERS_H_
 
+#include "Particles/WarpXParticleContainer.H"
 #include "Evolve/WarpXDtType.H"
 #include <AMReX_Config.H>
 #include <AMReX_REAL.H>
@@ -92,6 +93,20 @@ extern "C" {
     amrex::ParticleReal** warpx_getParticleArrays(int speciesnumber, int comp, int lev,
                                                   int* num_tiles, int** particles_per_tile);
 
+    amrex::ParticleReal** warpx_getParticleArraysFromCompName(
+        const char* char_species_name, const char* char_comp_name, int lev,
+        int* num_tiles, int** particles_per_tile);
+
+    amrex::ParticleReal** warpx_getParticleArraysUsingPC(
+        WarpXParticleContainer& myspc, int comp,
+        int lev, int* num_tiles, int** particles_per_tile);
+
+    int warpx_getParticleCompIndex(
+        const char* char_species_name, const char* char_comp_name);
+
+    void warpx_addRealComp(
+        const char* char_species_name, const char* char_comp_name, bool comm);
+
   void warpx_ComputeDt ();
   void warpx_MoveWindow (int step, bool move_j);