1 files changed, 132 insertions, 7 deletions

diff --git a/Source/Diagnostics/ParticleIO.cpp b/Source/Diagnostics/ParticleIO.cpp
index f15c084a0..f2a543ed5 100644
--- a/Source/Diagnostics/ParticleIO.cpp
+++ b/Source/Diagnostics/ParticleIO.cpp
@@ -5,6 +5,53 @@
 using namespace amrex;
 
 void
+RigidInjectedParticleContainer::ReadHeader (std::istream& is)
+{
+    is >> charge >> mass;
+    WarpX::GotoNextLine(is);
+
+    int nlevs;
+    is >> nlevs;
+    WarpX::GotoNextLine(is);
+
+    AMREX_ASSERT(zinject_plane_levels.size() == 0);
+    AMREX_ASSERT(done_injecting.size() == 0);
+
+    for (int i = 0; i < nlevs; ++i)
+    {
+        int zinject_plane_tmp;
+        is >> zinject_plane_tmp;
+        zinject_plane_levels.push_back(zinject_plane_tmp);
+        WarpX::GotoNextLine(is);        
+    } 
+
+    for (int i = 0; i < nlevs; ++i)
+    {
+        int done_injecting_tmp;
+        is >> done_injecting_tmp;
+        done_injecting.push_back(done_injecting_tmp);
+        WarpX::GotoNextLine(is);
+    }     
+}
+
+void
+RigidInjectedParticleContainer::WriteHeader (std::ostream& os) const
+{
+    // no need to write species_id
+    os << charge << " " << mass << "\n";
+    int nlevs = zinject_plane_levels.size();
+    os << nlevs << "\n";
+    for (int i = 0; i < nlevs; ++i)
+    {
+        os << zinject_plane_levels[i] << "\n";
+    }
+    for (int i = 0; i < nlevs; ++i)
+    {
+        os << done_injecting[i] << "\n";
+    }
+}
+
+void
 WarpXParticleContainer::ReadHeader (std::istream& is)
 {
     is >> charge >> mass;
@@ -27,17 +74,55 @@ MultiParticleContainer::Checkpoint (const std::string& dir) const
 }
 
 void
-MultiParticleContainer::WritePlotFile (const std::string& dir,
-                                       const Vector<int>& real_flags,
-                                       const Vector<std::string>& real_names) const
+MultiParticleContainer::WritePlotFile (const std::string& dir) const
 {
     Vector<std::string> int_names;    
     Vector<int> int_flags;
-    
+
     for (unsigned i = 0, n = species_names.size(); i < n; ++i) {
-	allcontainers[i]->WritePlotFile(dir, species_names[i],
-                                        real_flags, int_flags,
-                                        real_names, int_names);
+        auto& pc = allcontainers[i];                
+        if (pc->plot_species) {
+
+            Vector<std::string> real_names;
+            real_names.push_back("weight");
+
+            real_names.push_back("momentum_x");
+            real_names.push_back("momentum_y");
+            real_names.push_back("momentum_z");
+            
+            real_names.push_back("Ex");
+            real_names.push_back("Ey");
+            real_names.push_back("Ez");
+            
+            real_names.push_back("Bx");
+            real_names.push_back("By");
+            real_names.push_back("Bz");
+            
+#ifdef WARPX_RZ
+            real_names.push_back("theta");
+#endif
+            
+            if (WarpX::do_boosted_frame_diagnostic && pc->DoBoostedFrameDiags())
+            {
+                real_names.push_back("xold");
+                real_names.push_back("yold");
+                real_names.push_back("zold");
+                
+                real_names.push_back("uxold");
+                real_names.push_back("uyold");
+                real_names.push_back("uzold");
+            }
+                        
+            // Convert momentum to SI
+            pc->ConvertUnits(ConvertDirection::WarpX_to_SI);
+            // real_names contains a list of all particle attributes.
+            // pc->plot_flags is 1 or 0, whether quantity is dumped or not.
+            pc->WritePlotFile(dir, species_names[i],
+                              pc->plot_flags, int_flags,
+                              real_names, int_names);
+            // Convert momentum back to WarpX units
+            pc->ConvertUnits(ConvertDirection::SI_to_WarpX);
+        }
     }
 }
 
@@ -65,3 +150,43 @@ MultiParticleContainer::WriteHeader (std::ostream& os) const
     }
 }
 
+// Particle momentum is defined as gamma*velocity, which is neither 
+// SI mass*gamma*velocity nor normalized gamma*velocity/c.
+// This converts momentum to SI units (or vice-versa) to write SI data
+// to file.
+void
+PhysicalParticleContainer::ConvertUnits(ConvertDirection convert_direction)
+{
+    BL_PROFILE("PPC::ConvertUnits()");
+
+    // Compute conversion factor
+    Real factor = 1;
+    if (convert_direction == ConvertDirection::WarpX_to_SI){
+        factor = mass;
+    } else if (convert_direction == ConvertDirection::SI_to_WarpX){
+        factor = 1./mass;
+    }
+
+    for (int lev=0; lev<=finestLevel(); lev++){
+#ifdef _OPENMP
+#pragma omp parallel if (Gpu::notInLaunchRegion())
+#endif
+        for (WarpXParIter pti(*this, lev); pti.isValid(); ++pti)
+        {
+            // - momenta are stored as a struct of array, in `attribs`
+            auto& attribs = pti.GetAttribs();
+            Real* AMREX_RESTRICT ux = attribs[PIdx::ux].dataPtr();
+            Real* AMREX_RESTRICT uy = attribs[PIdx::uy].dataPtr();
+            Real* AMREX_RESTRICT uz = attribs[PIdx::uz].dataPtr();
+            // Loop over the particles and convert momentum
+            const long np = pti.numParticles();
+            ParallelFor( np,
+                [=] AMREX_GPU_DEVICE (long i) {
+                    ux[i] *= factor;
+                    uy[i] *= factor;
+                    uz[i] *= factor;
+                }
+            );
+        }
+    }
+}