add Lorentz Transform

Source/Diagnostics/ComputeDiagFunctors/BackTransformParticleFunctor.cpp

@@ -71,6 +71,111 @@ struct SelectParticles
 };
 
 
+struct LorentzTransformParticles
+{
+
+    using TmpParticles = WarpXParticleContainer::TmpParticles;
+
+    LorentzTransformParticles ( const WarpXParIter& a_pti, TmpParticles& tmp_particle_data,
+                                amrex::Real t_boost, amrex::Real dt,
+                                amrex::Real t_lab, int a_offset = 0)
+        : m_t_boost(t_boost), m_dt(dt), m_t_lab(t_lab)
+    {
+        using namespace amrex::literals;
+
+        if (tmp_particle_data.size() == 0) return;
+        m_get_position = GetParticlePosition(a_pti, a_offset);
+
+        auto& attribs = a_pti.GetAttribs();
+        m_wpnew = attribs[PIdx::w].dataPtr(); 
+        m_uxpnew = attribs[PIdx::ux].dataPtr(); 
+        m_uypnew = attribs[PIdx::ux].dataPtr(); 
+        m_uzpnew = attribs[PIdx::ux].dataPtr(); 
+
+        const auto lev = a_pti.GetLevel();
+        const auto index = a_pti.GetPairIndex();
+
+        m_xpold = tmp_particle_data[lev][index][TmpIdx::xold].dataPtr();
+        m_ypold = tmp_particle_data[lev][index][TmpIdx::yold].dataPtr();
+        m_zpold = tmp_particle_data[lev][index][TmpIdx::zold].dataPtr();
+        m_uxpold = tmp_particle_data[lev][index][TmpIdx::uxold].dataPtr();
+        m_uypold = tmp_particle_data[lev][index][TmpIdx::uyold].dataPtr();
+        m_uzpold = tmp_particle_data[lev][index][TmpIdx::uzold].dataPtr();
+
+        m_betaboost = WarpX::beta_boost;
+        m_gammaboost = WarpX::gamma_boost;
+        m_Phys_c = PhysConst::c;
+        m_inv_c2 = 1._rt/(m_Phys_c * m_Phys_c);
+        m_uzfrm = -m_gammaboost*m_betaboost*m_Phys_c;
+    }
+
+    template <typename DstData, typename SrcData>
+    AMREX_GPU_HOST_DEVICE
+    void operator () (DstData& dst, const SrcData src, int i_src, int i_dst) const noexcept
+    {
+        using namespace amrex::literals;
+        // get current src position
+        amrex::ParticleReal xpnew, ypnew, zpnew;
+        m_get_position(i_src, xpnew, ypnew, zpnew);
+
+        const amrex::Real gamma_new_p = std::sqrt(1.0_rt + m_inv_c2*
+                                        ( m_uxpnew[i_src] * m_uxpnew[i_src]
+                                        + m_uypnew[i_src] * m_uypnew[i_src]
+                                        + m_uzpnew[i_src] * m_uzpnew[i_src]));
+        const amrex::Real gamma_old_p = std::sqrt(1.0_rt + m_inv_c2*
+                                        ( m_uxpold[i_src] * m_uxpold[i_src]
+                                        + m_uypold[i_src] * m_uypold[i_src]
+                                        + m_uzpold[i_src] * m_uzpold[i_src]));
+        const amrex::Real t_new_p = m_gammaboost * m_t_boost - m_uzfrm * zpnew * m_inv_c2;
+        const amrex::Real z_new_p = m_gammaboost* ( zpnew + m_betaboost * m_Phys_c * m_t_boost);
+        const amrex::Real uz_new_p = m_gammaboost * m_uzpnew[i_src] - gamma_new_p * m_uzfrm;
+        const amrex::Real t_old_p = m_gammaboost * (m_t_boost - m_dt)
+                                    - m_uzfrm * m_zpold[i_src] * m_inv_c2;
+        const amrex::Real z_old_p = m_gammaboost * ( m_zpold[i_src] + m_betaboost
+                                                     * m_Phys_c * (m_t_boost - m_dt ) );
+        const amrex::Real uz_old_p = m_gammaboost * m_uzpold[i_src] - gamma_old_p * m_uzfrm;
+        // interpolate in time to t_lab
+        const amrex::Real weight_old = (t_new_p - m_t_lab)
+                                     / (t_new_p - t_old_p);
+        const amrex::Real weight_new = (m_t_lab - t_old_p)
+                                     / (t_new_p - t_old_p);
+        // weighted sum of old and new values
+        const amrex::ParticleReal xp = m_xpold[i_src] * weight_old + xpnew * weight_new;
+        const amrex::ParticleReal yp = m_ypold[i_src] * weight_old + ypnew * weight_new;
+        const amrex::ParticleReal zp = z_old_p * weight_old + z_new_p * weight_new;
+        const amrex::ParticleReal uxp = m_uxpold[i_src] * weight_old
+                                      + m_uxpnew[i_src] * weight_new;
+        const amrex::ParticleReal uyp = m_uypold[i_src] * weight_old
+                                      + m_uypnew[i_src] * weight_new;
+        const amrex::ParticleReal uzp = m_uzpold[i_src] * weight_old
+                                      + m_uzpnew[i_src] * weight_new;
+    }
+
+    GetParticlePosition m_get_position;
+
+    amrex::ParticleReal* AMREX_RESTRICT m_xpold = nullptr;
+    amrex::ParticleReal* AMREX_RESTRICT m_ypold = nullptr;
+    amrex::ParticleReal* AMREX_RESTRICT m_zpold = nullptr;
+
+    amrex::ParticleReal* AMREX_RESTRICT m_uxpold = nullptr;
+    amrex::ParticleReal* AMREX_RESTRICT m_uypold = nullptr;
+    amrex::ParticleReal* AMREX_RESTRICT m_uzpold = nullptr;
+
+    const amrex::ParticleReal* AMREX_RESTRICT m_uxpnew = nullptr;
+    const amrex::ParticleReal* AMREX_RESTRICT m_uypnew = nullptr;
+    const amrex::ParticleReal* AMREX_RESTRICT m_uzpnew = nullptr;
+    const amrex::ParticleReal* AMREX_RESTRICT m_wpnew = nullptr;
+
+    amrex::Real m_gammaboost;
+    amrex::Real m_betaboost;
+    amrex::Real m_Phys_c;
+    amrex::Real m_uzfrm;
+    amrex::Real m_inv_c2;
+    amrex::Real m_t_boost;
+    amrex::Real m_dt;
+    amrex::Real m_t_lab;
+};
+
 void
 BackTransformParticleFunctor::operator () (ParticleContainer& pc_dst, int i_buffer) const
 {