[do not merge] add diagnostic for laser fft rhs

src/diagnostics/Diagnostic.cpp

@@ -157,6 +157,8 @@ Diagnostic::Initialize (int nlev, bool use_laser) {
         geometry_name_to_level.emplace(geom_name, 0);
         all_comps_error_str << "Available components in base_geometry '" << geom_name << "':\n    ";
         geometry_name_to_output_comps[geom_name].insert(laser_io_name);
+        geometry_name_to_output_comps[geom_name].insert("rhs");
+        geometry_name_to_output_comps[geom_name].insert("rhs_fourier");
         all_comps_error_str << laser_io_name << "\n";
     }
     all_comps_error_str << "Additionally, 'all' and 'none' are supported as field_data\n"

src/diagnostics/OpenPMDWriter.cpp

@@ -148,7 +148,7 @@ OpenPMDWriter::WriteFieldData (
             std::vector< std::complex<double> > polarization {{1., 0.}, {0., 0.}};
             field.setAttribute("polarization", polarization);
             field_comp.storeChunkRaw(
-                reinterpret_cast<const std::complex<amrex::Real>*>(fd.m_F_laser.dataPtr()),
+                reinterpret_cast<const std::complex<amrex::Real>*>(fd.m_F_laser.dataPtr(icomp)),
                 chunk_offset, chunk_size);
         } else {
             field_comp.storeChunkRaw(fd.m_F.dataPtr(icomp), chunk_offset, chunk_size);

src/fields/Fields.cpp

@@ -520,10 +520,18 @@ Fields::Copy (const int current_N_level, const int i_slice, FieldDiagnosticData&
                 {
                     const amrex::Real x = i * dx + poff_diag_x;
                     const amrex::Real y = j * dy + poff_diag_y;
-                    diag_array_laser(i,j,k) += amrex::GpuComplex<amrex::Real> {
+                    diag_array_laser(i,j,k,0) += amrex::GpuComplex<amrex::Real> {
                         rel_z_data[k-k_min] * laser_array(x,y,WhichLaserSlice::n00j00_r),
                         rel_z_data[k-k_min] * laser_array(x,y,WhichLaserSlice::n00j00_i)
                     };
+                    diag_array_laser(i,j,k,1) += amrex::GpuComplex<amrex::Real> {
+                        rel_z_data[k-k_min] * laser_array(x,y,WhichLaserSlice::comp_rhs_r),
+                        rel_z_data[k-k_min] * laser_array(x,y,WhichLaserSlice::comp_rhs_i)
+                    };
+                    diag_array_laser(i,j,k,2) += amrex::GpuComplex<amrex::Real> {
+                        rel_z_data[k-k_min] * laser_array(x,y,WhichLaserSlice::comp_rhs_fourier_r),
+                        rel_z_data[k-k_min] * laser_array(x,y,WhichLaserSlice::comp_rhs_fourier_i)
+                    };
                 });
         }
     }

src/laser/MultiLaser.H

@@ -44,6 +44,10 @@ namespace WhichLaserSlice {
         np1jp2_i,
         chi,
         chi_initial,
+        comp_rhs_r,
+        comp_rhs_i,
+        comp_rhs_fourier_r,
+        comp_rhs_fourier_i,
         N
     };
 }

src/laser/MultiLaser.cpp

@@ -904,6 +904,8 @@ MultiLaser::AdvanceSliceMG (amrex::Real dt, int step)
                         rhs += arr(i, j, chi) * an00j00 * 2._rt;
                     }
                 }
+                arr(i, j, comp_rhs_r) = rhs.real();
+                arr(i, j, comp_rhs_i) = rhs.imag();
                 rhs_mg_arr(i,j,0) = rhs.real();
                 rhs_mg_arr(i,j,1) = rhs.imag();
             });
@@ -1068,6 +1070,8 @@ MultiLaser::AdvanceSliceFFT (const amrex::Real dt, int step)
                         - lapA
                         + ( -3._rt/(c*dt*dz) + 2._rt*I*djn/(c*dt) + 2._rt/(c*c*dt*dt) + I*2._rt*k0/(c*dt) ) * anm1j00;
                 }
+                arr(i, j, comp_rhs_r) = rhs.real();
+                arr(i, j, comp_rhs_i) = rhs.imag();
                 rhs_arr(i,j,0) = rhs;
             });
 
@@ -1085,11 +1089,14 @@ MultiLaser::AdvanceSliceFFT (const amrex::Real dt, int step)
         amrex::ParallelFor(
             to2D(bx),
             [=] AMREX_GPU_DEVICE(int i, int j) noexcept {
+                using namespace WhichLaserSlice;
                 // divide rhs_fourier by -(k^2+a)
                 amrex::Real kx = (i<imid) ? dkx*i : dkx*(i-Nx);
                 amrex::Real ky = (j<jmid) ? dky*j : dky*(j-Ny);
                 const Complex inv_k2a = abs(kx*kx + ky*ky + acoeff) > 0. ?
                     1._rt/(kx*kx + ky*ky + acoeff) : 0.;
+                arr(i, j, comp_rhs_fourier_r) = -inv_k2a.real();
+                arr(i, j, comp_rhs_fourier_i) = -inv_k2a.imag();
                 rhs_fourier_arr(i,j) *= -inv_k2a;
             });