Allow extra particle attributes (besides ux, uy, uz and w) to be set at particle creation in AddNParticles() (#2115)

* 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 functionality to AddNParticles() to allow extra particle attributes to also be set during particle creation

* 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

* added a test of the pywarpx bridge to get particle data and add new particle attributes at runtime

* changed all particle interacting functions in libwarpx to use the species name rather than id, also changed the functions to get particle array data to use the component name rather than index

* updated test according to PR #2119 changes

* removed unneeded BL_ASSERT(nattr == 1) statement

* fixed bug in add_particles to correctly determine the number of extra attributes

* fixed bug in AddNParticles if fewer attribute values are passed than the number of extra arrays for the species

* use isinstance(attr, ndarray) rather than type(attr) is np.ndarray

* generalize_runtime_comps_io

* fix OpenPMD

* fix OpenPMD

* fix plot flags in WritePlotFile

* fix offset and comment

* changed extra pid test to not use an underscore in the pid name

* switched _libwarpx.py::add_particles to use kwargs to accept the weight and extra attribute arrays

* License update in test file

Co-authored-by: Axel Huebl <[email protected]>

* fix typo

* added a test with unique_particles=False

* Apply suggestions from code review

Co-authored-by: Axel Huebl <[email protected]>

* updated docstring and comments

Co-authored-by: atmyers <[email protected]>
Co-authored-by: Axel Huebl <[email protected]>

Author: 3 people

Examples/Tests/ParticleDataPython/PICMI_inputs_2d.py

@@ -0,0 +1,157 @@
+from mpi4py import MPI
+from pywarpx import picmi
+import numpy as np
+
+##########################
+# MPI communicator setup
+##########################
+
+# split processor 0 into separate communicator from others
+comm_world = MPI.COMM_WORLD
+rank = comm_world.Get_rank()
+if rank == 0:
+    color = 0
+else:
+    color = 1
+new_comm = comm_world.Split(color)
+
+##########################
+# numerics parameters
+##########################
+
+dt = 7.5e-10
+
+# --- Nb time steps
+
+max_steps = 10
+
+# --- grid
+
+nx = 64
+ny = 64
+
+xmin = 0
+xmax = 0.03
+ymin = 0
+ymax = 0.03
+
+
+##########################
+# numerics components
+##########################
+
+grid = picmi.Cartesian2DGrid(
+    number_of_cells = [nx, ny],
+    lower_bound = [xmin, ymin],
+    upper_bound = [xmax, ymax],
+    lower_boundary_conditions = ['dirichlet', 'periodic'],
+    upper_boundary_conditions = ['dirichlet', 'periodic'],
+    lower_boundary_conditions_particles = ['absorbing', 'periodic'],
+    upper_boundary_conditions_particles = ['absorbing', 'periodic'],
+    moving_window_velocity = None,
+    warpx_max_grid_size = 32
+)
+
+solver = picmi.ElectrostaticSolver(
+    grid=grid, method='Multigrid', required_precision=1e-6,
+    warpx_self_fields_verbosity=0
+)
+
+##########################
+# physics components
+##########################
+
+electrons = picmi.Species(
+    particle_type='electron', name='electrons'
+)
+
+##########################
+# diagnostics
+##########################
+
+field_diag = picmi.FieldDiagnostic(
+    name = 'diag1',
+    grid = grid,
+    period = 10,
+    data_list = ['phi'],
+    write_dir = '.',
+    warpx_file_prefix = f'Python_particle_attr_access_plt_{color}'
+)
+
+##########################
+# simulation setup
+##########################
+
+sim = picmi.Simulation(
+    solver = solver,
+    time_step_size = dt,
+    max_steps = max_steps,
+    verbose = 1
+)
+
+sim.add_species(
+    electrons,
+    layout = picmi.GriddedLayout(
+        n_macroparticle_per_cell=[0, 0], grid=grid
+    )
+)
+sim.add_diagnostic(field_diag)
+
+sim.initialize_inputs()
+sim.initialize_warpx(mpi_comm=new_comm)
+
+##########################
+# python particle data access
+##########################
+
+from pywarpx import _libwarpx, callbacks
+
+_libwarpx.add_real_comp('electrons', 'newPid')
+
+def add_particles():
+
+    nps = 10
+    x = np.random.rand(nps) * 0.03
+    y = np.zeros(nps)
+    z = np.random.random(nps) * 0.03
+    ux = np.random.normal(loc=0, scale=1e3, size=nps)
+    uy = np.random.normal(loc=0, scale=1e3, size=nps)
+    uz = np.random.normal(loc=0, scale=1e3, size=nps)
+    w = np.ones(nps) * 2.0
+    newPid = 5.0
+
+    _libwarpx.add_particles(
+        species_name='electrons', x=x, y=y, z=z, ux=ux, uy=uy, uz=uz,
+        w=w, newPid=newPid, unique_particles=(not color)
+    )
+
+callbacks.installbeforestep(add_particles)
+
+##########################
+# simulation run
+##########################
+
+sim.step(max_steps - 1, mpi_comm=new_comm)
+
+##########################
+# check that the new PIDs are properly set
+##########################
+
+if color == 0:
+    assert (_libwarpx.get_particle_count('electrons') == 90)
+else:
+    assert (_libwarpx.get_particle_count('electrons') == 90)
+assert (_libwarpx.get_particle_comp_index('electrons', 'w') == 0)
+assert (_libwarpx.get_particle_comp_index('electrons', 'newPid') == 4)
+
+new_pid_vals = _libwarpx.get_particle_arrays(
+    'electrons', 'newPid', 0
+)
+for vals in new_pid_vals:
+    assert np.allclose(vals, 5)
+
+##########################
+# take the final sim step
+##########################
+
+sim.step(1)

Examples/Tests/ParticleDataPython/analysis.py

@@ -0,0 +1,18 @@
+#! /usr/bin/env python
+
+# Copyright 2021 Modern Electron
+#
+# License: BSD-3-Clause-LBNL
+
+# This script just checks that the PICMI file executed successfully.
+# If it did there will be a plotfile for the final step.
+
+import yt
+
+plotfile = 'Python_particle_attr_access_plt_000010'
+ds = yt.load( plotfile )  # noqa
+
+plotfile = 'Python_particle_attr_access_plt_100010'
+ds = yt.load( plotfile )  # noqa
+
+assert True

Python/pywarpx/_libwarpx.py

@@ -243,6 +243,15 @@ def get_nattr():
     # --- The -3 is because the comps include the velocites
     return libwarpx.warpx_nComps() - 3
 
+def get_nattr_species(species_name):
+    '''
+
+    Get the number of real attributes for the given species.
+
+    '''
+    return libwarpx.warpx_nCompsSpecies(
+        ctypes.c_char_p(species_name.encode('utf-8')))
+
 def amrex_init(argv, mpi_comm=None):
     # --- Construct the ctype list of strings to pass in
     argc = len(argv)
@@ -361,9 +370,8 @@ def getCellSize(direction, level=0):
 #
 #    libwarpx.warpx_ComputePMLFactors(lev, dt)
 
-def add_particles(species_name,
-                  x=0., y=0., z=0., ux=0., uy=0., uz=0., attr=0.,
-                  unique_particles=True):
+def add_particles(species_name, x=0., y=0., z=0., ux=0., uy=0., uz=0.,
+                  unique_particles=True, **kwargs):
     '''
 
     A function for adding particles to the WarpX simulation.
@@ -374,9 +382,11 @@ def add_particles(species_name,
     species_name     : the species to add the particle to
     x, y, z          : arrays or scalars of the particle positions (default = 0.)
     ux, uy, uz       : arrays or scalars of the particle momenta (default = 0.)
-    attr             : a 2D numpy array or scalar with the particle attributes (default = 0.)
     unique_particles : whether the particles are unique or duplicated on
                        several processes. (default = True)
+    kwargs           : dictionary containing an entry for the particle weights
+                       (with keyword 'w') and all the extra particle attribute
+                       arrays. If an attribute is not given it will be set to 0.
 
     '''
 
@@ -387,20 +397,20 @@ def add_particles(species_name,
     lenux = np.size(ux)
     lenuy = np.size(uy)
     lenuz = np.size(uz)
-    lenattr = np.size(attr)
 
     if (lenx == 0 or leny == 0 or lenz == 0 or lenux == 0 or
-        lenuy == 0 or lenuz == 0 or lenattr == 0):
+        lenuy == 0 or lenuz == 0):
         return
 
-    maxlen = max(lenx, leny, lenz, lenux, lenuy, lenuz, lenattr)
+    maxlen = max(lenx, leny, lenz, lenux, lenuy, lenuz)
     assert lenx==maxlen or lenx==1, "Length of x doesn't match len of others"
     assert leny==maxlen or leny==1, "Length of y doesn't match len of others"
     assert lenz==maxlen or lenz==1, "Length of z doesn't match len of others"
     assert lenux==maxlen or lenux==1, "Length of ux doesn't match len of others"
     assert lenuy==maxlen or lenuy==1, "Length of uy doesn't match len of others"
     assert lenuz==maxlen or lenuz==1, "Length of uz doesn't match len of others"
-    assert lenattr==maxlen or lenattr==1, "Length of attr doesn't match len of others"
+    for key, val in kwargs.items():
+        assert np.size(val)==1 or len(val)==maxlen, f"Length of {key} doesn't match len of others"
 
     if lenx == 1:
         x = np.array(x)*np.ones(maxlen)
@@ -414,15 +424,27 @@ def add_particles(species_name,
         uy = np.array(uy)*np.ones(maxlen)
     if lenuz == 1:
         uz = np.array(uz)*np.ones(maxlen,'d')
-    if lenattr == 1:
-        nattr = get_nattr()
-        attr = np.array(attr)*np.ones([maxlen,nattr])
+    for key, val in kwargs.items():
+        if np.size(val) == 1:
+            kwargs[key] = np.array(val)*np.ones(maxlen)
+
+    # --- The -3 is because the comps include the velocites
+    nattr = get_nattr_species(species_name) - 3
+    attr = np.zeros((maxlen, nattr))
+
+    for key, vals in kwargs.items():
+        if key == 'w':
+            attr[:,0] = vals
+        else:
+            # --- The -3 is because components 1 to 3 are velocities
+            attr[:,get_particle_comp_index(species_name, key)-3] = vals
 
     libwarpx.warpx_addNParticles(
         ctypes.c_char_p(species_name.encode('utf-8')), x.size,
-        x, y, z, ux, uy, uz, attr.shape[-1], attr, unique_particles
+        x, y, z, ux, uy, uz, nattr, attr, unique_particles
     )
 
+
 def get_particle_count(species_name):
     '''
 
@@ -444,6 +466,7 @@ def get_particle_count(species_name):
         ctypes.c_char_p(species_name.encode('utf-8'))
     )
 
+
 def get_particle_structs(species_name, level):
     '''
 

Regression/WarpX-tests.ini

@@ -2294,3 +2294,19 @@ compareParticles = 1
 particleTypes = electrons
 analysisRoutine = Examples/Modules/ParticleBoundaryProcess/analysis_absorption.py
 tolerance = 1.0e-4
+
+[Python_particle_attr_access]
+buildDir = .
+inputFile = Examples/Tests/ParticleDataPython/PICMI_inputs_2d.py
+runtime_params =
+customRunCmd = python PICMI_inputs_2d.py
+dim = 2
+addToCompileString = USE_PYTHON_MAIN=TRUE PYINSTALLOPTIONS="--user --prefix="
+restartTest = 0
+useMPI = 1
+numprocs = 3
+useOMP = 0
+numthreads = 0
+compileTest = 0
+doVis = 0
+analysisRoutine = Examples/Tests/ParticleDataPython/analysis.py

Source/Diagnostics/FlushFormats/FlushFormatPlotfile.cpp

@@ -374,7 +374,7 @@ FlushFormatPlotfile::WriteParticles(const std::string& dir,
         }, true);
 
         // real_names contains a list of all particle attributes.
-        // particle_diags[i].plot_flags is 1 or 0, whether quantity is dumped or not.
+        // real_flags & int_flags are 1 or 0, whether quantity is dumped or not.
         tmp.WritePlotFile(
             dir, particle_diags[i].getSpeciesName(),
             real_flags, int_flags,

Source/Particles/WarpXParticleContainer.cpp

@@ -139,13 +139,6 @@ WarpXParticleContainer::AddNParticles (int /*lev*/,
                                        const ParticleReal* vx, const ParticleReal* vy, const ParticleReal* vz,
                                        int nattr, const ParticleReal* attr, int uniqueparticles, amrex::Long id)
 {
-    // nattr is unused below but needed in the BL_ASSERT
-    amrex::ignore_unused(nattr);
-
-    BL_ASSERT(nattr == 1);
-
-    const ParticleReal* weight = attr;
-
     int ibegin, iend;
     if (uniqueparticles) {
         ibegin = 0;
@@ -175,6 +168,9 @@ WarpXParticleContainer::AddNParticles (int /*lev*/,
 
     std::size_t np = iend-ibegin;
 
+    // treat weight as a special attr since it will always be specified
+    Vector<ParticleReal> weight(np);
+
 #ifdef WARPX_DIM_RZ
     Vector<ParticleReal> theta(np);
 #endif
@@ -204,16 +200,15 @@ WarpXParticleContainer::AddNParticles (int /*lev*/,
         p.pos(1) = z[i];
 #endif
 
-        if ( (NumRuntimeRealComps()>0) || (NumRuntimeIntComps()>0) ){
-            DefineAndReturnParticleTile(0, 0, 0);
-        }
-
         pinned_tile.push_back(p);
+
+        // grab weight from the attr array
+        weight[i-ibegin] = attr[i*nattr];
     }
 
     if (np > 0)
     {
-        pinned_tile.push_back_real(PIdx::w , weight + ibegin, weight + iend);
+        pinned_tile.push_back_real(PIdx::w , weight.data(), weight.data() + np);
         pinned_tile.push_back_real(PIdx::ux,     vx + ibegin,     vx + iend);
         pinned_tile.push_back_real(PIdx::uy,     vy + ibegin,     vy + iend);
         pinned_tile.push_back_real(PIdx::uz,     vz + ibegin,     vz + iend);
@@ -236,9 +231,20 @@ WarpXParticleContainer::AddNParticles (int /*lev*/,
 #endif
         }
 
-        for (int i = PIdx::nattribs; i < NumRealComps(); ++i)
+        for (int j = PIdx::nattribs; j < NumRealComps(); ++j)
         {
-            pinned_tile.push_back_real(i, 0.0);
+            if (j - PIdx::nattribs < nattr - 1) {
+                // get the next attribute from attr array
+                Vector<ParticleReal> attr_vals(np);
+                for (int i = ibegin; i < iend; ++i)
+                {
+                    attr_vals[i-ibegin] = attr[j - PIdx::nattribs + 1 + i*nattr];
+                }
+                pinned_tile.push_back_real(j, attr_vals.data(), attr_vals.data() + np);
+            }
+            else {
+                pinned_tile.push_back_real(j, np, 0.0);
+            }
         }
 
         auto old_np = particle_tile.numParticles();

Source/Python/WarpXWrappers.cpp

@@ -132,6 +132,14 @@ extern "C"
         return PIdx::nattribs;
     }
 
+    int warpx_nCompsSpecies(const char* char_species_name)
+    {
+        auto & mypc = WarpX::GetInstance().GetPartContainer();
+        const std::string species_name(char_species_name);
+        auto & myspc = mypc.GetParticleContainerFromName(species_name);
+        return myspc.NumRealComps();
+    }
+
     int warpx_SpaceDim()
     {
         return AMREX_SPACEDIM;