From 4fa339e5951df30d06cd6084c5d7e47b7b3fa73c Mon Sep 17 00:00:00 2001
From: Roelof Groenewald <40245517+roelof-groenewald@users.noreply.github.com>
Date: Tue, 20 Jul 2021 17:55:42 -0700
Subject: 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
---
 Source/Python/WarpXWrappers.cpp | 50 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

(limited to 'Source/Python/WarpXWrappers.cpp')

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 ();
-- 
cgit v1.2.3