Reformatted BTD fields (#1158)

* fixing bug to initialize CellCenterFunctor for Bx

* create new BTB class and move FullDiagnostics functions from Diagnostics class

* [skip ci] Adding BTDiag class members and allocating them

* output_mf vector of buffers for both Full and BTD

* [skip ci] a few comments

* added more functions in BTDiagnostics for initialization

* [skip-ci] a small comment abour FlushRawData

* [skip-ci] Add new functor class for BTD

* Included BTDiagnostics class and skeleton of functions to be populated

* fix eol and make destructor of base Diag class virtual

* remove redundant query file_prefix in FullDiag class

* Apply suggestions from code review

adding PR suggestions

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>

* cleaning function calls and definitions as suggested in review

* uncomment the assert. It is valid for FullDiag and BTDiag

* Update Source/Diagnostics/Diagnostics.H

adding PR suggestion including comments

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>

* moved InitData to Diagnostics.cpp and call relevant functions from there

* moving FilterComputeAndPack to Diagnostics

* removing commented out code

* remove gpu vector containing map of fields from cc_mf to dst_mf. New diag machinery ensures cc_mf.ncomp == dst_mf.ncomp

* fix eol whitespace

* populating Init functions in BTDiagnostics

* DefineCellCenteredMultiFab should be called only once for all buffers

* add Init functions and constructing BackTransform functor

* call BackTransform operator which currently does nothing

* EOL

* average-down cc data to coarsest level

* loop over level used to allocate output buffer

* some comment about avg-down functions to be moved

* EOL

* Apply suggestions from code review

add pr suggestions

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>

* clarifying pure virtual and virtual functions

* removing namespace amrex and adding corresponding amrex:: prefix

* this should be nlev -- number of levels in the sim, not max allocated levels

* the number of levels to be output could be different from nlev, nmax_lev

* add missing semi-colon!

* add amrex:: prefix to MultiFab to comile psatd

* add new diag files to CMakeLists.txt

* fix compiler errors.

* add amrex:: prefix to fix psatd compilation error

* some more conflicts

* minor cleaning

* EOL

* flatten cc data using finestLevel() and use moving window dir

* EOL

* Apply suggestions from code review

PR suggestions

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

* making it readable and adding comments

* add missing semicolon

* add empty operator

* cc functor with ncomp=10 for LT. add prob domain. call functor prep data

* add cc_varnames for 10comp ccfunctor. define prob domain

* Add ith buffer to the operator for diag functors

* add ith buffer for diag functor; define prep functor data

* set plot rho to 1 for BTD

* adding prep data and LorentzTransform

* adding new BTD algorithm with some debugging print statements

* cleaning and ensuring use_gpu=true compiles

* EOL whitespace

* thanks bot! fixed redundant lev definition

* cleaning, adding comments, changing mlo and mhi in diags for boosted frame

* fixing compilation bugs

* ensuring accuracy in 2D by adding dim_map for diag lo and diag hi

* EOL fix

* prevent warning for unused variable

* some more comments and clean-up

* EOL

* fix bug with moving window conversion of mlo mhi

* restoring call to flush out diags after redistribute

* eol

* add reformatted input file for CI

* address PR review comments

* to prevent unused parameter warning

* temporarily removing CI test for reformatted slice BTD to fit CI time-limit

* call ComputePackFlush after syncing time

* Apply suggestions from code review

Co-authored-by: MaxThevenet <maxence.thevenet@desy.de>

* reverting DoComputeAndPack

Co-authored-by: MaxThevenet <mthevenet@lbl.gov>
Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>
Co-authored-by: MaxThevenet <maxence.thevenet@desy.de>

1 files changed, 189 insertions, 4 deletions

diff --git a/Source/Diagnostics/ComputeDiagFunctors/BackTransformFunctor.cpp b/Source/Diagnostics/ComputeDiagFunctors/BackTransformFunctor.cpp
index 2d8651efe..ed866e666 100644
--- a/Source/Diagnostics/ComputeDiagFunctors/BackTransformFunctor.cpp
+++ b/Source/Diagnostics/ComputeDiagFunctors/BackTransformFunctor.cpp
@@ -1,8 +1,193 @@
 #include "BackTransformFunctor.H"
+#include "WarpX.H"
+#include <AMReX_MultiFabUtil.H>
+#include <AMReX_MultiFabUtil_C.H>
+using namespace amrex;
 
-BackTransformFunctor::BackTransformFunctor(amrex::MultiFab const * mf_src, int lev,
-                                           const int ncomp, const amrex::IntVect crse_ratio)
-    : ComputeDiagFunctor(ncomp, crse_ratio), m_mf_src(mf_src), m_lev(lev)
+BackTransformFunctor::BackTransformFunctor (amrex::MultiFab const * mf_src, int lev,
+                                            const int ncomp, const int num_buffers,
+                                            amrex::Vector< std::string > varnames,
+                                            const amrex::IntVect crse_ratio)
+    : ComputeDiagFunctor(ncomp, crse_ratio), m_mf_src(mf_src), m_lev(lev), m_num_buffers(num_buffers), m_varnames(varnames)
 {
-    // Get the right slice of each field in the CC MultiFab, transform it and store it in the output.
+    InitData();
 }
+
+void
+BackTransformFunctor::operator ()(amrex::MultiFab& mf_dst, int dcomp, const int i_buffer) const
+{
+    // Perform back-transformation only if z slice is within the domain stored as 0/1
+    // in m_perform_backtransform[i_buffer]
+    if ( m_perform_backtransform[i_buffer] == 1) {
+        auto& warpx = WarpX::GetInstance();
+        auto geom = warpx.Geom(m_lev);
+        amrex::Real gamma_boost = warpx.gamma_boost;
+        int moving_window_dir = warpx.moving_window_dir;
+        amrex::Real beta_boost = std::sqrt( 1._rt - 1._rt/( gamma_boost * gamma_boost) );
+        bool interpolate = true;
+        std::unique_ptr< amrex::MultiFab > slice = nullptr;
+        int scomp = 0;
+        // Generate slice of the cell-centered multifab containing boosted-frame field-data
+        // at current z-boost location for the ith buffer
+        slice =  amrex::get_slice_data (moving_window_dir, m_current_z_boost[i_buffer],
+                     *m_mf_src, geom, scomp, m_mf_src->nComp(), interpolate);
+        // Perform in-place Lorentz-transform of all the fields stored in the slice.
+        LorentzTransformZ( *slice, gamma_boost, beta_boost);
+
+        // Create a 2D box for the slice in the boosted frame
+        amrex::Real dx = geom.CellSize(moving_window_dir);
+        // index corresponding to z_boost location in the boost-frame
+        int i_boost = static_cast<int> ( ( m_current_z_boost[i_buffer]
+                                            - geom.ProbLo(moving_window_dir) ) / dx );
+        // z-Slice at i_boost with x,y indices same as buffer_box
+        amrex::Box slice_box = m_buffer_box[i_buffer];
+        slice_box.setSmall(moving_window_dir, i_boost);
+        slice_box.setBig(moving_window_dir, i_boost);
+
+        // Make it a BoxArray
+        amrex::BoxArray slice_ba(slice_box);
+        slice_ba.maxSize( m_max_box_size );
+        // Define MultiFab with the distribution map of the destination multifab and
+        // containing all ten components that were in the slice generated from m_mf_src.
+        std::unique_ptr< amrex::MultiFab > tmp_slice_ptr = nullptr;
+        tmp_slice_ptr.reset( new MultiFab ( slice_ba, mf_dst.DistributionMap(),
+                                            slice->nComp(), 0) );
+        // Parallel copy the lab-frame data from "slice" MultiFab with
+        // ncomp=10 and boosted-frame dmap to "tmp_slice_ptr" MultiFab with
+        // ncomp=10 and dmap of the destination Multifab, which will store the final data
+        tmp_slice_ptr->copy( *slice, 0, 0, slice->nComp() );
+        // Now we will cherry pick only the user-defined fields from
+        // tmp_slice_ptr to dst_mf
+        const int k_lab = m_k_index_zlab[i_buffer];
+        const int ncomp_dst = mf_dst.nComp();
+        amrex::MultiFab& tmp = *tmp_slice_ptr;
+        for (amrex::MFIter mfi(tmp, TilingIfNotGPU()); mfi.isValid(); ++mfi)
+        {
+            // Box spanning the user-defined index-space for diagnostic, mf_dst
+            amrex::Box bx = m_buffer_box[i_buffer];
+            // Box of tmp_slice_ptr spanning full domain in x,y and z_boost
+            // of the respective buffer , i_buffer
+            const Box& tbx = mfi.tilebox();
+            // Modify bx, such that the z-index is equal to the sliced tmp_mf
+            bx.setSmall( moving_window_dir, tbx.smallEnd( moving_window_dir ) );
+            bx.setBig( moving_window_dir, tbx.bigEnd( moving_window_dir ) );
+            amrex::Array4<amrex::Real> src_arr = tmp[mfi].array();
+            amrex::Array4<amrex::Real> dst_arr = mf_dst[mfi].array();
+            const auto field_map_ptr = m_map_varnames.dataPtr();
+            amrex::ParallelFor( bx, ncomp_dst,
+                [=] AMREX_GPU_DEVICE(int i, int j, int k, int n)
+                {
+                    const int icomp = field_map_ptr[n];
+#if (AMREX_SPACEDIM == 3)
+                    dst_arr(i, j, k_lab, n) = src_arr(i, j, k, icomp);
+#else
+                    dst_arr(i, k_lab, k, n) = src_arr(i, j, k, icomp);
+#endif
+                } );
+        }
+
+        // Reset the temporary MultiFabs generated
+        slice.reset(new MultiFab);
+        slice.reset(nullptr);
+        tmp_slice_ptr.reset(new MultiFab);
+        tmp_slice_ptr.reset(nullptr);
+
+    }
+
+}
+
+void
+BackTransformFunctor::PrepareFunctorData (int i_buffer,
+                          bool ZSliceInDomain, amrex::Real current_z_boost,
+                          amrex::Box buffer_box, const int k_index_zlab )
+{
+    m_buffer_box[i_buffer] = buffer_box;
+    m_current_z_boost[i_buffer] = current_z_boost;
+    m_k_index_zlab[i_buffer] = k_index_zlab;
+    m_perform_backtransform[i_buffer] = 0;
+    if (ZSliceInDomain) m_perform_backtransform[i_buffer] = 1;
+}
+
+void
+BackTransformFunctor::InitData ()
+{
+
+    m_buffer_box.resize( m_num_buffers );
+    m_current_z_boost.resize( m_num_buffers );
+    m_perform_backtransform.resize( m_num_buffers );
+    m_k_index_zlab.resize( m_num_buffers );
+    m_map_varnames.resize( m_varnames.size() );
+
+    std::map<std::string, int> m_possible_fields_to_dump = {
+        {"Ex", 0},
+        {"Ey", 1},
+        {"Ez", 2},
+        {"Bx", 3},
+        {"By", 4},
+        {"Bz", 5},
+        {"jx", 6},
+        {"jy", 7},
+        {"jz", 8},
+        {"rho", 9}
+    };
+
+    for (int i = 0; i < m_varnames.size(); ++i)
+    {
+        m_map_varnames[i] = m_possible_fields_to_dump[ m_varnames[i] ] ;
+    }
+
+}
+
+void
+BackTransformFunctor::LorentzTransformZ (amrex::MultiFab& data, amrex::Real gamma_boost,
+                                         amrex::Real beta_boost) const
+{
+#ifdef _OPENMP
+#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
+#endif
+    for (amrex::MFIter mfi(data, TilingIfNotGPU()); mfi.isValid(); ++mfi) {
+        const amrex::Box& tbx = mfi.tilebox();
+        amrex::Array4< amrex::Real > arr = data[mfi].array();
+        amrex::Real clight = PhysConst::c;
+        amrex::Real inv_clight = 1.0_rt/clight;
+        // arr(x,y,z,comp) has ten-components namely,
+        // Ex Ey Ez Bx By Bz jx jy jz rho in that order.
+        amrex::ParallelFor( tbx,
+            [=] AMREX_GPU_DEVICE (int i, int j, int k)
+            {
+                // Back-transform the transverse electric and magnetic fields.
+                // Note that the z-components, Ez, Bz, are not changed by the transform.
+                amrex::Real e_lab, b_lab, j_lab, rho_lab;
+                // Transform Ex_boost (ncomp=0) & By_boost (ncomp=4) to lab-frame
+                e_lab = gamma_boost * ( arr(i, j, k, 0)
+                                        + beta_boost * clight * arr(i, j, k, 4) );
+                b_lab = gamma_boost * ( arr(i, j, k, 4)
+                                        + beta_boost * inv_clight * arr(i, j, k, 0) );
+                // Store lab-frame data in-place
+                arr(i, j, k, 0) = e_lab;
+                arr(i, j, k, 4) = b_lab;
+
+                // Transform Ey_boost (ncomp=1) & Bx_boost (ncomp=3) to lab-frame
+                e_lab = gamma_boost * ( arr(i, j, k, 1)
+                                        - beta_boost * clight * arr(i, j, k, 3) );
+                b_lab = gamma_boost * ( arr(i, j, k, 3)
+                                        - beta_boost * inv_clight * arr(i, j, k, 1) );
+                // Store lab-frame data in-place
+                arr(i, j, k, 1) = e_lab;
+                arr(i, j, k, 3) = b_lab;
+
+                // Transform charge density (ncomp=9)
+                // and z-component of current density (ncomp=8)
+                j_lab = gamma_boost * ( arr(i, j, k, 8)
+                                        + beta_boost * clight * arr(i, j, k, 9) );
+                rho_lab = gamma_boost * ( arr(i, j, k, 9)
+                                          + beta_boost * inv_clight * arr(i, j, k, 8) );
+                // Store lab-frame jz and rho in-place
+                arr(i, j, k, 8) = j_lab;
+                arr(i, j, k, 9) = rho_lab;
+            }
+        );
+    }
+
+}
+