adding custom output

Author: StaticObserver

pgens/BZjet/pgen.hpp

@@ -318,7 +318,7 @@ namespace user {
 
     const std::vector<real_t> xi_min;
     const std::vector<real_t> xi_max;
-    const real_t sigma0, sigma_max, inj_coeff, db_thr, temperature, m_eps;
+    const real_t sigma0, sigma_max, inj_coeff, db_thr, temperature, m_eps, inv_n0;
 
     InitFields<M, D>        init_flds;
     const Metadomain<S, M>* metadomain;
@@ -333,6 +333,7 @@ namespace user {
       , db_thr { p.template get<real_t>("setup.db_thr") }
       , temperature { p.template get<real_t>("setup.temperature") }
       , m_eps { p.template get<real_t>("setup.m_eps") }
+      , inv_n0 { ONE / p.template get<real_t>("scales.n0") }
       , init_flds { m.mesh().metric, m_eps }
       , metadomain { &m } {}
 
@@ -359,6 +360,110 @@ namespace user {
                                                          1.0,
                                                          true);
     }
+
+    void CustomFieldOutput(const std::string&   name,
+                           ndfield_t<M::Dim, 6> buffer,
+                           index_t              index,
+                           timestep_t,
+                           simtime_t,
+                           const Domain<S, M>&  domain) {
+
+       
+        if (name == "DB") {
+          if constexpr (M::Dim == Dim::_2D) {
+            const auto& EM = domain.fields.em;
+            const auto& metric = domain.mesh.metric;
+            Kokkos::parallel_for(
+            "DB",
+            domain.mesh.rangeActiveCells(),
+            Lambda(index_t i1, index_t i2) {
+              coord_t<M::Dim> xi { static_cast<real_t>(i1 - N_GHOSTS), static_cast<real_t>(i2 - N_GHOSTS) };
+              const vec_t<Dim::_3D> B_cntrv { EM(i1, i2, em::bx1),
+                                              EM(i1, i2, em::bx2),
+                                              EM(i1, i2, em::bx3) };
+              vec_t<Dim::_3D> B_cov { ZERO };
+              const vec_t<Dim::_3D> D_cntrv { EM(i1, i2, em::dx1),  
+                                              EM(i1, i2, em::dx2),
+                                              EM(i1, i2, em::dx3) };
+              metric.template transform<Idx::U, Idx::D>(xi, B_cntrv, B_cov);
+              buffer(i1, i2, index) = DOT(B_cov[0], B_cov[1], B_cov[2], D_cntrv[0], D_cntrv[1], D_cntrv[2])
+                                      / DOT(B_cntrv[0], B_cntrv[1], B_cntrv[2], B_cntrv[0], B_cntrv[1], B_cntrv[2]);
+            });
+          }
+        } else if (name == "fRho"){
+          auto T3 = domain.fields.buff;
+          auto& mesh         = domain.mesh;
+          const auto use_weights = params.template get<bool>(
+            "particles.use_weights");
+          const auto ni2 = mesh.n_active(in::x2);
+          std::vector<unsigned short> specs {};
+          for (auto& sp : domain.species) {
+            if (sp.mass() > 0) {
+              specs.push_back(sp.index());
+            }
+          }
+          const auto& metric = domain.mesh.metric;
+
+          // compute T_00, T_01, T_11
+          auto scatter_buff = Kokkos::Experimental::create_scatter_view(T3);
+          for (const auto& sp : specs) {
+            auto& prtl_spec = domain.species[sp - 1];
+            // clang-format off
+            Kokkos::parallel_for(
+              "ComputeMoments",
+              prtl_spec.rangeActiveParticles(),
+              kernel::ParticleMoments_kernel<S, M, FldsID::T, 3>({0, 0}, scatter_buff, 0u,
+                                                                  prtl_spec.i1, prtl_spec.i2, prtl_spec.i3,
+                                                                  prtl_spec.dx1, prtl_spec.dx2, prtl_spec.dx3,
+                                                                  prtl_spec.ux1, prtl_spec.ux2, prtl_spec.ux3,
+                                                                  prtl_spec.phi, prtl_spec.weight, prtl_spec.tag,
+                                                                  prtl_spec.mass(), prtl_spec.charge(),
+                                                                  use_weights,
+                                                                  metric, mesh.flds_bc(),
+                                                                  ni2, inv_n0, ZERO));
+            Kokkos::parallel_for(
+              "ComputeMoments",
+              prtl_spec.rangeActiveParticles(),
+              kernel::ParticleMoments_kernel<S, M, FldsID::T, 3>({0, 1}, scatter_buff, 1u,
+                                                                  prtl_spec.i1, prtl_spec.i2, prtl_spec.i3,
+                                                                  prtl_spec.dx1, prtl_spec.dx2, prtl_spec.dx3,
+                                                                  prtl_spec.ux1, prtl_spec.ux2, prtl_spec.ux3,
+                                                                  prtl_spec.phi, prtl_spec.weight, prtl_spec.tag,
+                                                                  prtl_spec.mass(), prtl_spec.charge(),
+                                                                  use_weights,
+                                                                  metric, mesh.flds_bc(),
+                                                                  ni2, inv_n0, ZERO));
+            Kokkos::parallel_for(
+              "ComputeMoments",
+              prtl_spec.rangeActiveParticles(),
+              kernel::ParticleMoments_kernel<S, M, FldsID::T, 3>({1, 1}, scatter_buff, 2u,
+                                                                  prtl_spec.i1, prtl_spec.i2, prtl_spec.i3,
+                                                                  prtl_spec.dx1, prtl_spec.dx2, prtl_spec.dx3,
+                                                                  prtl_spec.ux1, prtl_spec.ux2, prtl_spec.ux3,
+                                                                  prtl_spec.phi, prtl_spec.weight, prtl_spec.tag,
+                                                                  prtl_spec.mass(), prtl_spec.charge(),
+                                                                  use_weights,
+                                                                  metric, mesh.flds_bc(),
+                                                                  ni2, inv_n0, ZERO));
+            // clang-format on
+          }
+          Kokkos::Experimental::contribute(T3, scatter_buff);
+
+          if constexpr (M::Dim == Dim::_2D) {
+            Kokkos::parallel_for(
+              "fRho",
+              domain.mesh.rangeActiveCells(),
+              Lambda(index_t i1, index_t i2) {
+                coord_t<M::Dim> xi { static_cast<real_t>(i1 - N_GHOSTS), static_cast<real_t>(i2 - N_GHOSTS) };
+                real_t alpha_ij = metric.alpha(xi);
+                real_t beta_ij = metric.beta1(xi);
+                buffer(i1, i2, index) = (T3(i1, i2, 0) - TWO * beta_ij * T3(i1, i2, 1) + SQR(beta_ij) * T3(i1, i2, 2)) / SQR(alpha_ij);
+              });
+          }
+        } else {
+          raise::Error("Custom output not provided", HERE);
+        }
+      }
   };
 
 } // namespace user