Merge pull request #43 from germasch/pr/particle-tag

particle-tag: add support for particle tags
psc-code · Jul 27, 2019 · 0b84050 · 0b84050
2 parents 80ff6ab + f23dc6b
commit 0b84050
Show file tree

Hide file tree

Showing 14 changed files with 56 additions and 15 deletions.
diff --git a/src/include/const_accessor_simple.hxx b/src/include/const_accessor_simple.hxx
@@ -67,6 +67,7 @@ struct ConstParticleProxySimple
   real_t m()  const { return mprts_.grid().kinds[kind()].m; }
   int kind()  const { return prt_.kind; }
   psc::particle::Id id()    const { return prt_.id(); }
+  psc::particle::Tag tag() const { return prt_.tag(); }
 
   Double3 position() const
   {

diff --git a/src/include/injector_simple.hxx b/src/include/injector_simple.hxx
@@ -29,7 +29,7 @@ struct InjectorSimple
       auto prt = Particle{{real_t(new_prt.x[0] - patch.xb[0]), real_t(new_prt.x[1] - patch.xb[1]), real_t(new_prt.x[2] - patch.xb[2])},
 			  {real_t(new_prt.u[0]), real_t(new_prt.u[1]), real_t(new_prt.u[2])},
 			  real_t(new_prt.w * mprts_.grid().kinds[new_prt.kind].q),
-			  new_prt.kind, mprts_.uid_gen()};
+			  new_prt.kind, mprts_.uid_gen(), new_prt.tag};
       mprts_[p_].push_back(prt);
     }
 

diff --git a/src/include/particle.h b/src/include/particle.h
@@ -11,6 +11,7 @@ namespace particle
 {
 
 using Id = uint64_t;
+using Tag = int;
 
 // ======================================================================
 // Inject
@@ -23,14 +24,15 @@ struct Inject
   using Real = double;
   using Real3 = Vec3<Real>;
 
-  Inject(const Real3& x, const Real3& u, Real w, int kind)
-    : x{x}, u{u}, w{w}, kind{kind}
+  Inject(const Real3& x, const Real3& u, Real w, int kind, Tag tag = {})
+    : x{x}, u{u}, w{w}, kind{kind}, tag{tag}
   {}
 
   Real3 x;
   Real3 u;
   Real w;
   int kind;
+  Tag tag;
 };
 
 } // namespace particle

diff --git a/src/include/particle_simple.hxx b/src/include/particle_simple.hxx
@@ -15,7 +15,7 @@ struct ParticleSimple
 
   ParticleSimple() = default;
 
-  KG_INLINE ParticleSimple(Real3 x, Real3 u, real_t qni_wni, int kind, psc::particle::Id id)
+  KG_INLINE ParticleSimple(Real3 x, Real3 u, real_t qni_wni, int kind, psc::particle::Id id, psc::particle::Tag tag)
     : x{x},
       u{u},
       kind{kind},
@@ -31,6 +31,7 @@ struct ParticleSimple
   KG_INLINE bool operator!=(const ParticleSimple& other) const { return !(*this == other); }
 
   KG_INLINE psc::particle::Id id()  const { return 0; }
+  KG_INLINE psc::particle::Tag tag()  const { return 0; }
 
 public:
   Real3 x;

diff --git a/src/include/particle_with_id.h b/src/include/particle_with_id.h
@@ -16,8 +16,8 @@ struct ParticleWithId
   ParticleWithId() = default;
 
   KG_INLINE ParticleWithId(Real3 x, Real3 u, real_t qni_wni, int kind,
-                           psc::particle::Id id)
-    : x{x}, u{u}, kind{kind}, qni_wni{qni_wni}, id_{id}
+                           psc::particle::Id id, psc::particle::Tag tag)
+    : x{x}, u{u}, kind{kind}, qni_wni{qni_wni}, id_{id}, tag_{tag}
   {}
 
   KG_INLINE bool operator==(const ParticleWithId& other) const
@@ -32,13 +32,15 @@ struct ParticleWithId
   }
 
   KG_INLINE psc::particle::Id id() const { return id_; }
+  KG_INLINE psc::particle::Tag tag() const { return tag_; }
 
 public:
   Real3 x;
   Real3 u;
   int kind;
   real_t qni_wni;
   psc::particle::Id id_;
+  psc::particle::Tag tag_;
 };
 
 template <typename R>
@@ -59,5 +61,6 @@ class ForComponents<ParticleWithId<R>>
     func("kind", [](Particle& prt) { return &prt.kind; });
     func("qni_wni", [](Particle& prt) { return &prt.qni_wni; });
     func("id", [](Particle& prt) { return &prt.id_; });
+    func("tag", [](Particle& prt) { return &prt.tag_; });
   }
 };
diff --git a/src/include/setup_particles.hxx b/src/include/setup_particles.hxx
@@ -7,6 +7,7 @@ struct psc_particle_npt
   double n;    ///< density
   double p[3]; ///< momentum
   double T[3]; ///< temperature
+  psc::particle::Tag tag;
 };
 
 // ======================================================================
@@ -81,7 +82,7 @@ struct SetupParticles
       }
     }
 
-    return psc::particle::Inject{pos, {pxi, pyi, pzi}, wni, npt.kind};
+    return psc::particle::Inject{pos, {pxi, pyi, pzi}, wni, npt.kind, npt.tag};
   }
 
   // ----------------------------------------------------------------------
@@ -111,7 +112,6 @@ struct SetupParticles
     }
 
     auto inj = mprts.injector();
-
     for (int p = 0; p < mprts.n_patches(); ++p) {
       auto ldims = grid.ldims;
       auto injector = inj[p];

diff --git a/src/libpsc/cuda/cuda_mparticles.cu b/src/libpsc/cuda/cuda_mparticles.cu
@@ -518,7 +518,7 @@ std::vector<typename cuda_mparticles<BS>::Particle> cuda_mparticles<BS>::get_par
     int kind = cuda_float_as_int(xi4[n].w);
     prts.emplace_back(Real3{xi4[n].x, xi4[n].y, xi4[n].z},
 	              Real3{pxi4[n].x, pxi4[n].y, pxi4[n].z},
-	              pxi4[n].w, kind, 0);
+	              pxi4[n].w, kind, psc::particle::Id(), psc::particle::Tag());
 
 #if 0
     uint b = blockIndex(xi4[n], p);

diff --git a/src/libpsc/cuda/cuda_mparticles.cuh b/src/libpsc/cuda/cuda_mparticles.cuh
@@ -69,7 +69,7 @@ struct ParticleCudaStorage
   __host__ __device__
   operator DParticleCuda()
   {
-    return {{xi4.x, xi4.y, xi4.z}, {pxi4.x, pxi4.y, pxi4.z}, pxi4.w, cuda_float_as_int(xi4.w), 0};
+    return {{xi4.x, xi4.y, xi4.z}, {pxi4.x, pxi4.y, pxi4.z}, pxi4.w, cuda_float_as_int(xi4.w), psc::particle::Id(), psc::particle::Tag()};
   }
 
   float4 xi4;

diff --git a/src/libpsc/cuda/injector_buffered.hxx b/src/libpsc/cuda/injector_buffered.hxx
@@ -38,7 +38,7 @@ struct InjectorBuffered
       auto x = Double3::fromPointer(new_prt.x) - patch.xb;
       auto u = Double3::fromPointer(new_prt.u);
       real_t q = mprts.grid().kinds[new_prt.kind].q;
-      raw(Particle{Real3(x), Real3(u), q * real_t(new_prt.w), new_prt.kind, mprts.uid_gen()});
+      raw(Particle{Real3(x), Real3(u), q * real_t(new_prt.w), new_prt.kind, mprts.uid_gen(), new_prt.tag});
     }
 
     // FIXME do we want to keep this? or just have a psc::particle::Inject version instead?

diff --git a/src/libpsc/cuda/psc_particles_cuda.cxx b/src/libpsc/cuda/psc_particles_cuda.cxx
@@ -24,7 +24,7 @@ static void copy_from(MparticlesBase& mprts_base, MparticlesBase& mprts_other_ba
     for (auto prt: accessor[p]) {
       using real_t = typename MparticlesCuda::real_t;
       using Real3 = typename MparticlesCuda::Real3;
-      injector.raw({Real3(prt.x()), Real3(prt.u()), real_t(prt.qni_wni()), prt.kind(), prt.id()});
+      injector.raw({Real3(prt.x()), Real3(prt.u()), real_t(prt.qni_wni()), prt.kind(), prt.id(), prt.tag()});
     }
   }
 }
@@ -43,7 +43,7 @@ static void copy_to(MparticlesBase& mprts_base, MparticlesBase& mprts_other_base
     for (auto prt: accessor[p]) {
       using real_t = typename MP::real_t;
       using Real3 = typename MP::Real3;
-      mprts_other[p].push_back({Real3(prt.x()), Real3(prt.u()), real_t(prt.qni_wni()), prt.kind(), prt.id()});
+      mprts_other[p].push_back({Real3(prt.x()), Real3(prt.u()), real_t(prt.qni_wni()), prt.kind(), prt.id(), prt.tag()});
     }
   }
 }

diff --git a/src/libpsc/psc_output_particles/output_particles_hdf5_impl.hxx b/src/libpsc/psc_output_particles/output_particles_hdf5_impl.hxx
@@ -14,6 +14,7 @@ struct hdf5_prt {
   float px, py, pz;
   float q, m, w;
   psc::particle::Id id;
+  psc::particle::Tag tag;
 };
 
 #define H5_CHK(ierr) assert(ierr >= 0)
@@ -70,6 +71,7 @@ struct OutputParticlesHdf5 : OutputParticlesParams, OutputParticlesBase
     H5Tinsert(id, "m" , HOFFSET(struct hdf5_prt, m) , H5T_NATIVE_FLOAT);
     H5Tinsert(id, "w" , HOFFSET(struct hdf5_prt, w) , H5T_NATIVE_FLOAT);
     H5Tinsert(id, "id" , HOFFSET(struct hdf5_prt, id) , ToHdf5Type<psc::particle::Id>::H5Type());
+    H5Tinsert(id, "tag" , HOFFSET(struct hdf5_prt, tag) , ToHdf5Type<psc::particle::Tag>::H5Type());
     prt_type = id;
 
     // set hi to gdims by default (if not set differently before)
@@ -246,6 +248,7 @@ struct OutputParticlesHdf5 : OutputParticlesParams, OutputParticlesBase
 		  arr[nn].m  = prt.m();
 		  arr[nn].w  = prt.w();
 		  arr[nn].id  = prt.id();
+		  arr[nn].tag  = prt.tag();
 		}
 	      }
 	    }

diff --git a/src/libpsc/psc_particles_impl.cxx b/src/libpsc/psc_particles_impl.cxx
@@ -17,7 +17,7 @@ struct Convert
     using Real3 = typename MparticlesTo::Real3;
 
     return {Real3(prt_from.x), Real3(prt_from.u),
-	real_t(prt_from.qni_wni), prt_from.kind, prt_from.id()};
+	real_t(prt_from.qni_wni), prt_from.kind, prt_from.id(), prt_from.tag()};
   }
 };
 

diff --git a/src/libpsc/tests/test_mparticles.cxx b/src/libpsc/tests/test_mparticles.cxx
@@ -521,6 +521,37 @@ TEST(TestSetupParticles, Id)
   }
 }
 
+TEST(TestSetupParticles, Tag)
+{
+  using Mparticles = MparticlesSimple<ParticleWithId<double>>;
+
+  auto domain = Grid_t::Domain{{1, 2, 2},
+			       {10., 20., 20.}, {},
+			       {1, 1, 1}};
+  auto kinds = Grid_t::Kinds{{1., 100., "i"}, {-1., 1., "e"}};
+  auto prm = Grid_t::NormalizationParams::dimensionless();
+  prm.nicell = 2;
+  Grid_t grid{domain, {}, kinds, {prm}, .1};
+  Mparticles mprts{grid};
+
+  SetupParticles<Mparticles> setup_particles;
+  std::vector<uint> n_prts_by_patch;
+  setup_particles.setup_particles(mprts, n_prts_by_patch,
+				  [&](int kind, Double3 crd, psc_particle_npt& npt) {
+      npt.n = 1;
+      npt.tag = psc::particle::Tag{kind * 10};
+    });
+
+  auto n_cells = grid.domain.gdims[0] * grid.domain.gdims[1] * grid.domain.gdims[2];
+  EXPECT_EQ(mprts.size(), n_cells * kinds.size() * prm.nicell);
+
+  for (int p = 0; p < mprts.n_patches(); p++) {
+    for (auto &prt : mprts[p]) {
+      EXPECT_EQ(prt.tag(), prt.kind * 10);
+    }
+  }
+}
+
 int main(int argc, char** argv)
 {
   MPI_Init(&argc, &argv);

diff --git a/src/libpsc/vpic/psc_particles_vpic.cxx b/src/libpsc/vpic/psc_particles_vpic.cxx
@@ -49,7 +49,7 @@ void copy_to(MparticlesBase& mprts_from_base, MparticlesBase& mprts_to_base)
 	auto u = Vec3<float>{vprt.ux, vprt.uy, vprt.uz};
 	auto kind = sp->id;
 	auto qni_wni = float(vprt.w * dVi) * float(mprts_to.grid().kinds[kind].q);
-	mprts_to[p].push_back({x, u, qni_wni, kind, 0});
+	mprts_to[p].push_back({x, u, qni_wni, kind, psc::particle::Id{}, psc::particle::Tag{}});
       }
     }
   }