Merge pull request #384 from JamesMcClung/pr/polymorphic-marder

Polymorphic Marder

Author: JamesMcClung

src/include/gauss_corrector_base.hxx

@@ -0,0 +1,7 @@
+#pragma once
+
+template <typename MfieldsState, typename Mparticles>
+struct GaussCorrectorBase
+{
+  virtual void correct_gauss(MfieldsState& mflds, Mparticles& mprts) = 0;
+};

src/include/psc.hxx

@@ -8,6 +8,7 @@
 #include <setup_particles.hxx>
 
 #include "../libpsc/vpic/fields_item_vpic.hxx"
+#include "gauss_corrector_base.hxx"
 #include "diagnostic_base.hxx"
 #include "injector_base.hxx"
 #include "external_current_base.hxx"
@@ -111,13 +112,13 @@ struct Psc
   using Sort = typename PscConfig::Sort;
   using Collision = typename PscConfig::Collision;
   using Checks = typename PscConfig::Checks;
-  using Marder = typename PscConfig::Marder;
   using PushParticles = typename PscConfig::PushParticles;
   using PushFields = typename PscConfig::PushFields;
   using Bnd = typename PscConfig::Bnd;
   using BndFields = typename PscConfig::BndFields;
   using BndParticles = typename PscConfig::BndParticles;
   using Dim = typename PscConfig::Dim;
+  using GaussCorrectorBaseT = GaussCorrectorBase<MfieldsState, Mparticles>;
   using DiagnosticBaseT = DiagnosticBase<Mparticles, MfieldsState>;
   using ParticleDiagnosticBaseT = ParticleDiagnosticBase<Mparticles>;
   using InjectorBaseT = InjectorBase<Mparticles, MfieldsState>;
@@ -127,16 +128,14 @@ struct Psc
   // ctor
 
   Psc(const PscParams& params, Grid_t& grid, MfieldsState& mflds,
-      Mparticles& mprts, Balance& balance, Collision& collision, Checks& checks,
-      Marder& marder)
+      Mparticles& mprts, Balance& balance, Collision& collision, Checks& checks)
     : p_{params},
       grid_{&grid},
       mflds_{mflds},
       mprts_{mprts},
       balance_{balance},
       collision_{collision},
       checks_{checks},
-      marder_{marder},
       bndp_{grid},
       checkpointing_{params.write_checkpoint_every_step}
   {
@@ -165,6 +164,13 @@ struct Psc
   // TODO: improve ownership model: we should own these objects (i.e., use
   // unique_ptr), but don't want to burden the user with C++ boilerplate.
 
+  void add_gauss_corrector(GaussCorrectorBaseT* corrector)
+  {
+    if (corrector) {
+      gauss_correctors_.push_back(corrector);
+    }
+  }
+
   void add_injector(InjectorBaseT* injector)
   {
     assert(injector);
@@ -429,6 +435,15 @@ struct Psc
     bndf.fill_ghosts_E(mflds_);
     bnd_.fill_ghosts(mflds_, EX, EX + 3);
     prof_stop(pr_bndf);
+
+    if (p_.marder_interval > 0 && timestep % p_.marder_interval == 0) {
+      mpi_printf(comm, "***** Performing Marder correction...\n");
+      prof_start(pr_marder);
+      for (auto corrector : gauss_correctors_) {
+        corrector->correct_gauss(mflds_, mprts_);
+      }
+      prof_stop(pr_marder);
+    }
     // state is now: x^{n+3/2}, p^{n+1}, E^{n+3/2}, B^{n+1}
 
     // === field propagation B^{n+1} -> B^{n+3/2}
@@ -451,16 +466,6 @@ struct Psc
       prof_stop(pr_checks);
     }
 
-    // E at t^{n+3/2}, particles at t^{n+3/2}
-    // B at t^{n+3/2} (Note: that is not its natural time,
-    // but div B should be == 0 at any time...)
-    if (p_.marder_interval > 0 && timestep % p_.marder_interval == 0) {
-      mpi_printf(comm, "***** Performing Marder correction...\n");
-      prof_start(pr_marder);
-      marder_(mflds_, mprts_);
-      prof_stop(pr_marder);
-    }
-
     if (checks_.gauss.should_do_check(timestep)) {
       mpi_printf(comm, "***** Checking gauss...\n");
       prof_restart(pr_checks);
@@ -532,7 +537,7 @@ protected:
   Balance& balance_;
   Collision& collision_;
   Checks& checks_;
-  Marder& marder_;
+  std::vector<GaussCorrectorBaseT*> gauss_correctors_;
   std::vector<DiagnosticBaseT*> diagnostics_;
   std::vector<InjectorBaseT*> injectors_;
   std::vector<ExternalCurrentBaseT*> external_currents_;
@@ -559,14 +564,13 @@ protected:
 // makePscIntegrator
 
 template <typename PscConfig, typename MfieldsState, typename Mparticles,
-          typename Balance, typename Collision, typename Checks,
-          typename Marder>
+          typename Balance, typename Collision, typename Checks>
 Psc<PscConfig> makePscIntegrator(const PscParams& params, Grid_t& grid,
                                  MfieldsState& mflds, Mparticles& mprts,
                                  Balance& balance, Collision& collision,
-                                 Checks& checks, Marder& marder)
+                                 Checks& checks)
 {
-  return {params, grid, mflds, mprts, balance, collision, checks, marder};
+  return {params, grid, mflds, mprts, balance, collision, checks};
 }
 
 // ======================================================================

src/libpsc/psc_push_fields/marder_impl.hxx

@@ -1,6 +1,7 @@
 
 #pragma once
 
+#include "gauss_corrector_base.hxx"
 #include "fields.hxx"
 #include "writer_mrc.hxx"
 #include "mpi_dtype_traits.hxx"
@@ -66,45 +67,35 @@ inline void correct(const Grid_t& grid, E1& efield, const Int3& efield_ib,
   assert(efield_ib == -grid.ibn);
   assert(mf_ib == -grid.ibn);
 
-  Real3 fac = {.5f * grid.dt * diffusion / grid.domain.dx[0],
-               .5f * grid.dt * diffusion / grid.domain.dx[1],
-               .5f * grid.dt * diffusion / grid.domain.dx[2]};
+  Real3 fac = .5f * real_t(grid.dt) * diffusion * Real3(grid.domain.dx_inv);
 
   for (int p = 0; p < grid.n_patches(); p++) {
     Int3 lx, rx, ly, ry, lz, rz;
     detail::find_limits(grid, p, lx, rx, ly, ry, lz, rz);
 
-    if (!grid.isInvar(0)) {
-      Int3 l = lx, r = rx;
-      auto ex = efield.view(_all, _all, _all, 0, p);
-      auto res = mf.view(_all, _all, _all, 0, p);
-      ex.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2])) =
-        ex.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2])) +
-        (res.view(_s(l[0] + 1, r[0] + 1), _s(l[1], r[1]), _s(l[2], r[2])) -
-         res.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2]))) *
-          fac[0];
-    }
+    Int3 ls[3] = {lx, ly, lz};
+    Int3 rs[3] = {rx, ry, rz};
 
-    {
-      Int3 l = ly, r = ry;
-      auto ey = efield.view(_all, _all, _all, 1, p);
-      auto res = mf.view(_all, _all, _all, 0, p);
-      ey.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2])) =
-        ey.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2])) +
-        (res.view(_s(l[0], r[0]), _s(l[1] + 1, r[1] + 1), _s(l[2], r[2])) -
-         res.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2]))) *
-          fac[1];
-    }
+    auto res = mf.view(_all, _all, _all, 0, p);
+    for (int d = 0; d < 3; d++) {
+      if (grid.isInvar(d)) {
+        continue;
+      }
+
+      Int3 l = ls[d];
+      Int3 r = rs[d];
+      auto e_comp = efield.view(_all, _all, _all, d, p);
+
+      gt::gslice s1x = _s(l[0], r[0]);
+      gt::gslice s1y = _s(l[1], r[1]);
+      gt::gslice s1z = _s(l[2], r[2]);
+
+      gt::gslice s2[3] = {_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2])};
+      s2[d] = _s(l[d] + 1, r[d] + 1);
 
-    {
-      Int3 l = lz, r = rz;
-      auto ez = efield.view(_all, _all, _all, 2, p);
-      auto res = mf.view(_all, _all, _all, 0, p);
-      ez.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2])) =
-        ez.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2])) +
-        (res.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2] + 1, r[2] + 1)) -
-         res.view(_s(l[0], r[0]), _s(l[1], r[1]), _s(l[2], r[2]))) *
-          fac[2];
+      e_comp.view(s1x, s1y, s1z) =
+        e_comp.view(s1x, s1y, s1z) +
+        (res.view(s2[0], s2[1], s2[2]) - res.view(s1x, s1y, s1z)) * fac[d];
     }
   }
 }
@@ -198,12 +189,14 @@ inline void correct(const Grid_t& grid, E1& efield, const Int3& efield_ib,
 } // namespace marder
 } // namespace psc
 
-template <typename S, typename D, typename ITEM_RHO, typename BND>
-class MarderCommon
+template <typename MFIELDS_STATE, typename MPARTICLES, typename ITEM_RHO,
+          typename BND>
+class MarderCommon : public GaussCorrectorBase<MFIELDS_STATE, MPARTICLES>
 {
 public:
-  using storage_type = S;
-  using dim_t = D;
+  using MfieldsState = MFIELDS_STATE;
+  using Mparticles = MPARTICLES;
+  using storage_type = typename MfieldsState::Storage;
   using Item_rho_t = ITEM_RHO;
   using Bnd = BND;
   using real_t = typename storage_type::value_type;
@@ -241,23 +234,16 @@ public:
     }
   }
 
-  // ----------------------------------------------------------------------
-  // operator()
-  //
-  // Do the modified marder correction (See eq.(5, 7, 9, 10) in Mardahl and
-  // Verboncoeur, CPC, 1997)
-
-  template <typename Mparticles>
-  void operator()(const Grid_t& grid, storage_type& mflds, const Int3& mflds_ib,
-                  Mparticles& mprts)
+  void correct_gauss(MfieldsState& mflds, Mparticles& mprts) override
   {
-    auto efield = mflds.view(_all, _all, _all, _s(EX, EX + 3), _all);
-    auto efield_ib = mflds_ib;
+    const Grid_t& grid = mflds.grid();
+    auto efield = mflds.storage().view(_all, _all, _all, _s(EX, EX + 3), _all);
+    Int3 efield_ib = mflds.ib();
 
     double inv_sum = 0.;
     for (int d = 0; d < 3; d++) {
       if (!grid.isInvar(d)) {
-        inv_sum += 1. / sqr(grid.domain.dx[d]);
+        inv_sum += sqr(grid.domain.dx_inv[d]);
       }
     }
     double diffusion_max = 1. / 2. / (.5 * grid.dt) / inv_sum;
@@ -270,7 +256,7 @@ public:
     // and expected to be filled when done, so we're playing it safe for the
     // time being.
     for (int i = 0; i < loop_; i++) {
-      bnd_.fill_ghosts(grid, mflds, mflds_ib, EX, EX + 3);
+      bnd_.fill_ghosts(mflds, EX, EX + 3);
       auto dive = psc::mflds::interior(grid, psc::item::div_nc(grid, efield));
 
       Int3 res_ib = -grid.ibn;
@@ -282,21 +268,8 @@ public:
 
       psc::marder::correct(grid, efield, efield_ib, res, res_ib, diffusion);
     }
-    bnd_.fill_ghosts(grid, mflds, mflds_ib, EX, EX + 3);
-  }
-
-  template <typename MfieldsState, typename Mparticles>
-  void operator()(MfieldsState& mflds, Mparticles& mprts)
-  {
-    static int pr;
-    if (!pr) {
-      pr = prof_register("marder", 1., 0, 0);
-    }
-
-    prof_start(pr);
-    (*this)(mprts.grid(), mflds.storage(), mflds.ib(), mprts);
-    prof_stop(pr);
-  }
+    bnd_.fill_ghosts(mflds, EX, EX + 3);
+  };
 
   // private:
   real_t diffusion_; //< diffusion coefficient for Marder correction
@@ -306,16 +279,18 @@ public:
   WriterMRC io_; //< for debug dumping
 };
 
-template <typename S, typename D>
-using Marder_ = MarderCommon<S, D, Moment_rho_1st_nc<S, D>, Bnd_>;
+template <typename MFIELDS_STATE, typename MPARTICLES, typename D>
+using Marder_ =
+  MarderCommon<MFIELDS_STATE, MPARTICLES,
+               Moment_rho_1st_nc<typename MFIELDS_STATE::Storage, D>, Bnd_>;
 
 #ifdef USE_CUDA
 
 #include "psc_particles_single.h"
 #include "mparticles_cuda.hxx"
 #include "fields_item_moments_1st_cuda.hxx"
 
-template <typename D>
-using MarderCuda = MarderCommon<MfieldsStateCuda::Storage, D,
+template <typename Mparticles, typename D>
+using MarderCuda = MarderCommon<MfieldsStateCuda, Mparticles,
                                 Moment_rho_1st_nc_cuda<D>, BndCuda3>;
 #endif

src/libpsc/tests/test_boundary_injector.cxx

@@ -38,7 +38,6 @@ using Mparticles = PscConfig::Mparticles;
 using Balance = PscConfig::Balance;
 using Collision = PscConfig::Collision;
 using Checks = PscConfig::Checks;
-using Marder = PscConfig::Marder;
 using OutputParticles = PscConfig::OutputParticles;
 
 Grid_t* setupGrid()
@@ -133,10 +132,9 @@ TEST(BoundaryInjectorTest, Integration1Particle)
 
   Balance balance{.1};
   Collision collision{grid, 0, 0.1};
-  Marder marder(grid, 0.9, 3, false);
 
   auto psc = makePscIntegrator<PscConfig>(psc_params, grid, mflds, mprts,
-                                          balance, collision, checks, marder);
+                                          balance, collision, checks);
 
   psc.add_injector(
     new BoundaryInjector<ParticleGenerator, typename PscConfig::PushParticles>(
@@ -191,10 +189,9 @@ TEST(BoundaryInjectorTest, IntegrationManyParticles)
 
   Balance balance{.1};
   Collision collision{grid, 0, 0.1};
-  Marder marder(grid, 0.9, 3, false);
 
   auto psc = makePscIntegrator<PscConfig>(psc_params, grid, mflds, mprts,
-                                          balance, collision, checks, marder);
+                                          balance, collision, checks);
 
   psc.add_injector(
     new BoundaryInjector<ParticleGenerator, PscConfig::PushParticles>(
@@ -249,7 +246,6 @@ TEST(BoundaryInjectorTest, IntegrationManySpecies)
 
   Balance balance{.1};
   Collision collision{grid, 0, 0.1};
-  Marder marder(grid, 0.9, 3, false);
 
   auto inject_electrons =
     BoundaryInjector<ParticleGenerator, PscConfig::PushParticles>{
@@ -259,7 +255,7 @@ TEST(BoundaryInjectorTest, IntegrationManySpecies)
       ParticleGenerator(-1, 1), grid};
 
   auto psc = makePscIntegrator<PscConfig>(psc_params, grid, mflds, mprts,
-                                          balance, collision, checks, marder);
+                                          balance, collision, checks);
 
   psc.add_injector(&inject_ions);
   psc.add_injector(&inject_electrons);

src/libpsc/tests/test_open_bcs_integration.cxx

@@ -26,7 +26,6 @@ using Mparticles = PscConfig::Mparticles;
 using Balance = PscConfig::Balance;
 using Collision = PscConfig::Collision;
 using Checks = PscConfig::Checks;
-using Marder = PscConfig::Marder;
 using OutputParticles = PscConfig::OutputParticles;
 
 // ======================================================================
@@ -104,10 +103,9 @@ TEST(OpenBcsTest, IntegrationY)
 
   Balance balance{.1};
   Collision collision{grid, 0, 0.1};
-  Marder marder(grid, 0.9, 3, false);
 
   auto psc = makePscIntegrator<PscConfig>(psc_params, *grid_ptr, mflds, mprts,
-                                          balance, collision, checks, marder);
+                                          balance, collision, checks);
 
   // ----------------------------------------------------------------------
   // set up initial conditions

src/libpsc/tests/test_reflective_bcs.cxx

@@ -26,7 +26,6 @@ using Mparticles = PscConfig::Mparticles;
 using Balance = PscConfig::Balance;
 using Collision = PscConfig::Collision;
 using Checks = PscConfig::Checks;
-using Marder = PscConfig::Marder;
 using OutputParticles = PscConfig::OutputParticles;
 
 // ======================================================================