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

3 files changed, 156 insertions, 0 deletions

diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py
index f1cc6564d..ca298768d 100755
--- a/Python/pywarpx/_libwarpx.py
+++ b/Python/pywarpx/_libwarpx.py
@@ -145,6 +145,8 @@ libwarpx.amrex_init.argtypes = (ctypes.c_int, _LP_LP_c_char)
 libwarpx.amrex_init_with_inited_mpi.argtypes = (ctypes.c_int, _LP_LP_c_char, _MPI_Comm_type)
 libwarpx.warpx_getParticleStructs.restype = _LP_particle_p
 libwarpx.warpx_getParticleArrays.restype = _LP_LP_c_particlereal
+libwarpx.warpx_getParticleArraysFromCompName.restype = _LP_LP_c_particlereal
+libwarpx.warpx_getParticleCompIndex.restype = ctypes.c_int
 libwarpx.warpx_getEfield.restype = _LP_LP_c_real
 libwarpx.warpx_getEfieldLoVects.restype = _LP_c_int
 libwarpx.warpx_getEfieldCP.restype = _LP_LP_c_real
@@ -503,6 +505,51 @@ def get_particle_arrays(species_number, comp, level):
     return particle_data
 
 
+def get_particle_arrays_from_comp_name(species_name, comp_name, level):
+    '''
+
+    This returns a list of numpy arrays containing the particle array data
+    on each tile for this process.
+
+    The data for the numpy arrays are not copied, but share the underlying
+    memory buffer with WarpX. The numpy arrays are fully writeable.
+
+    Parameters
+    ----------
+
+        species_name   : the species name that the data will be returned for
+        comp_name      : the component of the array data that will be returned.
+
+    Returns
+    -------
+
+        A List of numpy arrays.
+
+    '''
+
+    particles_per_tile = _LP_c_int()
+    num_tiles = ctypes.c_int(0)
+    data = libwarpx.warpx_getParticleArraysFromCompName(
+        ctypes.c_char_p(species_name.encode('utf-8')),
+        ctypes.c_char_p(comp_name.encode('utf-8')),
+        level, ctypes.byref(num_tiles), ctypes.byref(particles_per_tile)
+    )
+
+    particle_data = []
+    for i in range(num_tiles.value):
+        arr = np.ctypeslib.as_array(data[i], (particles_per_tile[i],))
+        try:
+            # This fails on some versions of numpy
+            arr.setflags(write=1)
+        except ValueError:
+            pass
+        particle_data.append(arr)
+
+    _libc.free(particles_per_tile)
+    _libc.free(data)
+    return particle_data
+
+
 def get_particle_x(species_number, level=0):
     '''
 
@@ -643,6 +690,50 @@ def get_particle_theta(species_number, level=0):
         raise Exception('get_particle_r: There is no theta coordinate with 2D Cartesian')
 
 
+def get_particle_comp_index(species_name, pid_name):
+    '''
+
+    Get the component index for a given particle attribute. This is useful
+    to get the corrent ordering of attributes when adding new particles using
+    `add_particles()`.
+
+    Parameters
+    ----------
+
+        species_name   : the species name that the data will be returned for
+        pid_name       : string that is used to identify the new component
+
+    Returns
+    -------
+
+        Integer corresponding to the index of the requested attribute
+
+    '''
+    return libwarpx.warpx_getParticleCompIndex(
+        ctypes.c_char_p(species_name.encode('utf-8')),
+        ctypes.c_char_p(pid_name.encode('utf-8'))
+    )
+
+
+def add_real_comp(species_name, pid_name, comm=True):
+    '''
+
+    Add a real component to the particle data array.
+
+    Parameters
+    ----------
+
+        species_name   : the species name for which the new component will be added
+        pid_name       : string that is used to identify the new component
+        comm           : should the component be communicated
+
+    '''
+    libwarpx.warpx_addRealComp(
+        ctypes.c_char_p(species_name.encode('utf-8')),
+        ctypes.c_char_p(pid_name.encode('utf-8')), comm
+    )
+
+
 def _get_mesh_field_list(warpx_func, level, direction, include_ghosts):
     """
      Generic routine to fetch the list of field data arrays.
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);