BTD-RZ Add multiple modes

Examples/Tests/BTD_rz/analysis_BTD_laser_antenna.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+
+# Copyright 2022
+# Authors: Revathi Jambunathan, Remi Lehe
+#
+# This tests checks the backtransformed diagnostics by emitting a laser
+# (with the antenna) in the boosted-frame and then checking that the
+# fields recorded by the backtransformed diagnostics have the right amplitude,
+# wavelength, and envelope (i.e. gaussian envelope with the right duration.
+
+import numpy as np
+from openpmd_viewer import OpenPMDTimeSeries
+from scipy.constants import c, e, m_e
+from scipy.optimize import curve_fit
+
+
+def gaussian_laser( z, a0, z0_phase, z0_prop, ctau, lambda0 ):
+    """
+    Returns a Gaussian laser profile
+    """
+    k0 = 2*np.pi/lambda0
+    E0 = a0*m_e*c**2*k0/e
+    return( E0*np.exp( - (z-z0_prop)**2/ctau**2 ) \
+                *np.cos( k0*(z-z0_phase) ) )
+
+# Fit the on-axis profile to extract the phase (a.k.a. CEP)
+def fit_function(z, z0_phase):
+    return( gaussian_laser( z, a0, z0_phase,
+                            z0_b+Lprop_b, ctau0, lambda0 ) )
+
+# The values must be consistent with the values provided in the simulation input
+t_current = 80e-15   # Time of the snapshot1
+c = 299792458;
+z0_antenna = -1.e-6  # position of laser
+lambda0 = 0.8e-6     # wavelength of the signal
+tau0 = 10e-15        # duration of the signal
+ctau0 = tau0 * c
+a0 = 15              # amplitude
+t_peak = 20e-15      # Time at which laser reaches its peak
+Lprop_b = c*t_current
+z0_b = z0_antenna - c * t_peak
+
+ts = OpenPMDTimeSeries('./diags/back_rz')
+Ex, info = ts.get_field('E', 'x', iteration=1, slice_across='r')
+
+fit_result = curve_fit( fit_function, info.z, Ex,
+                        p0=np.array([z0_b+Lprop_b]) )
+z0_fit = fit_result[0]
+
+Ex_fit = gaussian_laser( info.z, a0, z0_fit, z0_b+Lprop_b, ctau0, lambda0)
+
+## Check that the a0 agrees within 5% of the predicted value
+assert np.allclose( Ex, Ex_fit, atol=0.18*Ex.max() )

Examples/Tests/BTD_rz/inputs_rz_z_boosted_BTD

@@ -0,0 +1,64 @@
+# Maximum number of time steps
+warpx.zmax_plasma_to_compute_max_step = 500e-6
+# number of grid points
+amr.n_cell = 32 256
+
+# Maximum allowable size of each subdomain in the problem domain;
+#    this is used to decompose the domain for parallel calculations.
+amr.max_grid_size = 128
+
+# Maximum level in hierarchy (for now must be 0, i.e., one level in total)
+amr.max_level = 0
+
+# Geometry
+geometry.dims = RZ
+geometry.prob_lo     =  0.e-6 -20.e-6
+geometry.prob_hi     =  40.e-6  0.e-6
+
+boundary.field_lo = none absorbing_silver_mueller
+boundary.field_hi = absorbing_silver_mueller absorbing_silver_mueller
+
+# Boosted frame and moving window
+warpx.do_moving_window = 1
+warpx.moving_window_dir = z
+warpx.moving_window_v = 1.0 # in units of the speed of light
+warpx.gamma_boost = 10.
+warpx.boost_direction = z
+
+
+# Verbosity
+warpx.verbose = 1
+warpx.n_rz_azimuthal_modes = 2
+
+# Algorithms
+warpx.cfl = 1.0
+warpx.use_filter = 0
+
+
+# Order of particle shape factors
+algo.particle_shape = 1
+
+# Laser
+lasers.names        = laser1
+laser1.profile      = Gaussian
+laser1.position     = 0. 0. -1.e-6 # This point is on the laser plane
+laser1.direction    = 0. 0. 1.     # The plane normal direction
+laser1.polarization = 1. 0. 0.     # The main polarization vector
+laser1.a0           = 1.5e1        # Maximum amplitude of the laser field
+laser1.profile_waist = 10.e-6      # The waist of the laser (in meters)
+laser1.profile_duration = 10.e-15  # The duration of the laser (in seconds)
+laser1.profile_t_peak = 20.e-15    # The time at which the laser reaches its peak (in seconds)
+laser1.profile_focal_distance = 1.e-6  # Focal distance from the antenna (in meters)
+laser1.wavelength = 0.8e-6         # The wavelength of the laser (in meters)
+
+# Diagnostics
+diagnostics.diags_names = diag1 back_rz
+diag1.intervals = 50
+diag1.diag_type = Full
+
+back_rz.diag_type = BackTransformed
+back_rz.dt_snapshots_lab = 80.e-15
+back_rz.fields_to_plot = Er Et Ez Br Bt Bz jr jt jz rho
+back_rz.format = openpmd
+back_rz.buffer_size = 32
+back_rz.num_snapshots_lab = 2

Regression/WarpX-tests.ini

@@ -3715,3 +3715,19 @@ doVis = 0
 compareParticles = 1
 particleTypes = electron ion
 analysisRoutine = Examples/Tests/VayDeposition/analysis.py
+
+[BTD_rz]
+buildDir = .
+inputFile = Examples/Tests/BTD_rz/inputs_rz_z_boosted_BTD
+runtime_params =
+dim = 2
+addToCompileString = USE_RZ=TRUE
+cmakeSetupOpts = -DWarpX_DIMS=RZ
+restartTest = 0
+useMPI = 1
+numprocs = 2
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+analysisRoutine = Examples/Tests/BTD_rz/analysis_BTD_laser_antenna.py

Source/Diagnostics/BTDiagnostics.cpp

@@ -198,7 +198,6 @@ BTDiagnostics::ReadParameters ()
         m_crse_ratio == amrex::IntVect(1),
         "Only support for coarsening ratio of 1 in all directions is included for BTD\n"
         );
-    WARPX_ALWAYS_ASSERT_WITH_MESSAGE(WarpX::n_rz_azimuthal_modes==1, "Currently only one mode is supported for BTD");
     // Read list of back-transform diag parameters requested by the user //
     amrex::ParmParse pp_diag_name(m_diag_name);
 
@@ -481,8 +480,12 @@ BTDiagnostics::DefineCellCenteredMultiFab(int lev)
     ba.coarsen(m_crse_ratio);
     amrex::DistributionMapping dmap = warpx.DistributionMap(lev);
     int ngrow = 1;
+#ifdef WARPX_DIM_RZ
+    int ncomps = WarpX::ncomps * static_cast<int>(m_cellcenter_varnames.size());
+#else
     int ncomps = static_cast<int>(m_cellcenter_varnames.size());
-    m_cell_centered_data[lev] = std::make_unique<amrex::MultiFab>(ba, dmap, ncomps, ngrow);
+#endif
+    WarpX::AllocInitMultiFab(m_cell_centered_data[lev], ba, dmap, ncomps, amrex::IntVect(ngrow), "cellcentered_BTD",0._rt);
 
 }
 
@@ -520,7 +523,7 @@ BTDiagnostics::InitializeFieldFunctors (int lev)
         int nvars = static_cast<int>(m_varnames.size());
         m_all_field_functors[lev][i] = std::make_unique<BackTransformFunctor>(
                   m_cell_centered_data[lev].get(), lev,
-                  nvars, m_num_buffers, m_varnames);
+                  nvars, m_num_buffers, m_varnames, m_varnames_fields);
     }
 
     // Define all cell-centered functors required to compute cell-centere data
@@ -632,7 +635,8 @@ BTDiagnostics::InitializeFieldFunctorsRZopenPMD (int lev)
         int nvars = static_cast<int>(m_varnames.size());
         m_all_field_functors[lev][i] = std::make_unique<BackTransformFunctor>(
                                        m_cell_centered_data[lev].get(), lev,
-                                       nvars, m_num_buffers, m_varnames);
+                                       nvars, m_num_buffers, m_varnames,
+                                       m_varnames_fields);
     }
 
     // Reset field functors for cell-center multifab

Source/Diagnostics/ComputeDiagFunctors/BackTransformFunctor.H

@@ -47,6 +47,7 @@ public:
     BackTransformFunctor ( const amrex::MultiFab * const mf_src, const int lev,
                            const int ncomp, const int num_buffers,
                            amrex::Vector< std::string > varnames,
+                           amrex::Vector< std::string > varnames_fields,
                            const amrex::IntVect crse_ratio= amrex::IntVect(1));
 
     /** \brief Lorentz-transform mf_src for the ith buffer and write the result in mf_dst.
@@ -118,6 +119,8 @@ private:
     amrex::Vector<int> m_k_index_zlab;
     /** Vector of user-defined field names to be stored in the output multifab */
     amrex::Vector< std::string > m_varnames;
+    /** Vector of user-defined field names without modifications for rz modes */
+    amrex::Vector< std::string > m_varnames_fields;
 
     /** max grid size used to generate BoxArray to define output MultiFabs */
     int m_max_box_size = 256;

Source/Diagnostics/ComputeDiagFunctors/BackTransformFunctor.cpp

@@ -34,8 +34,10 @@ using namespace amrex;
 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)
+                                            amrex::Vector< std::string > varnames_fields,
+                                            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), m_varnames_fields(varnames_fields)
 {
     InitData();
 }
@@ -112,14 +114,32 @@ BackTransformFunctor::operator ()(amrex::MultiFab& mf_dst, int /*dcomp*/, const
             const Box& tbx = mfi.tilebox();
             amrex::Array4<amrex::Real> src_arr = tmp[mfi].array();
             amrex::Array4<amrex::Real> dst_arr = mf_dst[mfi].array();
+#ifdef WARPX_DIM_RZ
+            const int n_rz_comp = WarpX::ncomps;
+#endif
             amrex::ParallelFor( tbx, ncomp_dst,
                 [=] AMREX_GPU_DEVICE(int i, int j, int k, int n)
                 {
+                    // Field id that corresponds to the nth user-requested component
                     const int icomp = field_map_ptr[n];
 #if defined(WARPX_DIM_3D)
                     dst_arr(i, j, k_lab, n) = src_arr(i, j, k, icomp);
-#else
+#elif defined(WARPX_DIM_XZ)
                     dst_arr(i, k_lab, k, n) = src_arr(i, j, k, icomp);
+#elif defined(WARPX_DIM_RZ)
+                    // rzcomp below gives the component id, 0 to (n_rz_comp-1) for a given field
+                    const int rzcomp = n % n_rz_comp;
+                    // Accessing the correct rz component from the cell-centered multifab
+                    // that has back-transformed fields and storing it for the appropriate user-requested field, icomp
+                    // For example, for 2 rz modes, we have three components (n_rz_comp=3) for each field
+                    // If n = 4 gives icomp = 1 (for Et) obtained from field_map_ptr,
+                    //           rzcomp = 4 - int(floor(4/3))*3 = 4 - 3 = 1
+                    // Thus we are accessing real component of mode 1 of Et (note that modes go from 0 to 1)
+                    // Since the fields are stored contiguously in src_arr, icomp*n_rz_comp + rz_comp accesses
+                    // real part of mode 1 for Et (1*3+1) = 4
+                    dst_arr(i, k_lab, k, n) = src_arr(i, j, k, icomp*n_rz_comp+rzcomp);
+#else
+                    dst_arr(k_lab, j, k, n) = src_arr(i, j, k, icomp);
 #endif
                 } );
         }
@@ -185,7 +205,12 @@ BackTransformFunctor::InitData ()
 
     for (int i = 0; i < m_varnames.size(); ++i)
     {
+#ifdef WARPX_DIM_RZ
+        const int field_id = i / WarpX::ncomps;
+        m_map_varnames[i] = m_possible_fields_to_dump[ m_varnames_fields[field_id] ];
+#else
         m_map_varnames[i] = m_possible_fields_to_dump[ m_varnames[i] ] ;
+#endif
     }
 
 }
@@ -202,6 +227,46 @@ BackTransformFunctor::LorentzTransformZ (amrex::MultiFab& data, amrex::Real gamm
         amrex::Array4< amrex::Real > arr = data[mfi].array();
         amrex::Real clight = PhysConst::c;
         amrex::Real inv_clight = 1.0_rt/clight;
+#ifdef WARPX_DIM_RZ
+        const int n_rcomps = WarpX::ncomps;
+        amrex::ParallelFor( tbx,
+            [=] AMREX_GPU_DEVICE (int i, int j, int k)
+            {
+                for (int mode_comp = 0; mode_comp < n_rcomps; ++mode_comp) {
+                    // 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 Er_boost & Bt_boost to lab-frame for corresponding mode (mode_comp)
+                    e_lab = gamma_boost * ( arr(i, j, k, n_rcomps*0 + mode_comp)
+                                            + beta_boost * clight * arr(i, j, k, n_rcomps*4+ mode_comp) );
+                    b_lab = gamma_boost * ( arr(i, j, k, n_rcomps*4 + mode_comp)
+                                            + beta_boost * inv_clight * arr(i, j, k, n_rcomps*0 + mode_comp) );
+                    // Store lab-frame data in-place
+                    arr(i, j, k, n_rcomps*0 + mode_comp) = e_lab;
+                    arr(i, j, k, n_rcomps*4 + mode_comp) = b_lab;
+
+                    // Transform Et_boost & Br_boost to lab-frame for corresponding mode (mode_comp)
+                    e_lab = gamma_boost * ( arr(i, j, k, n_rcomps*1 + mode_comp)
+                                            - beta_boost * clight * arr(i, j, k, n_rcomps*3 + mode_comp) );
+                    b_lab = gamma_boost * ( arr(i, j, k, n_rcomps*3 + mode_comp)
+                                            - beta_boost * inv_clight * arr(i, j, k, n_rcomps*1 + mode_comp) );
+                    // Store lab-frame data in-place
+                    arr(i, j, k, n_rcomps*1 + mode_comp) = e_lab;
+                    arr(i, j, k, n_rcomps*3 + mode_comp) = b_lab;
+
+                    // Transform charge density z-component of current density
+                    j_lab = gamma_boost * ( arr(i, j, k, n_rcomps*8 + mode_comp)
+                                            + beta_boost * clight * arr(i, j, k, n_rcomps*9 + mode_comp) );
+                    rho_lab = gamma_boost * ( arr(i, j, k, n_rcomps*9 + mode_comp)
+                                              + beta_boost * inv_clight * arr(i, j, k, n_rcomps*8 + mode_comp) );
+                    // Store lab-frame jz and rho in-place
+                    arr(i, j, k, n_rcomps*8 + mode_comp) = j_lab;
+                    arr(i, j, k, n_rcomps*9 + mode_comp) = rho_lab;
+                }
+            }
+        );
+#else
         // 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,
@@ -239,6 +304,7 @@ BackTransformFunctor::LorentzTransformZ (amrex::MultiFab& data, amrex::Real gamm
                 arr(i, j, k, 9) = rho_lab;
             }
         );
+#endif
     }
 
 }

Source/Diagnostics/MultiDiagnostics.cpp

@@ -25,9 +25,6 @@ MultiDiagnostics::MultiDiagnostics ()
             alldiags[i] = std::make_unique<FullDiagnostics>(i, diags_names[i]);
         } else if ( diags_types[i] == DiagTypes::BackTransformed ){
             alldiags[i] = std::make_unique<BTDiagnostics>(i, diags_names[i]);
-#ifdef WARPX_DIM_RZ
-            ablastr::warn_manager::WMRecordWarning("MultiDiagnostics", "BackTransformed diagnostics for fields is not yet fully implemented in RZ. Field output might be incorrect.");
-#endif
         } else if ( diags_types[i] == DiagTypes::BoundaryScraping ){
             alldiags[i] = std::make_unique<BoundaryScrapingDiagnostics>(i, diags_names[i]);
         } else {