add diagnostic for laser fft rhs

Hi-PACE · Sep 26, 2024 · ba8d067 · ba8d067
1 parent a6be0ce
commit ba8d067
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 1 deletion.
diff --git a/src/diagnostics/Diagnostic.cpp b/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"

diff --git a/src/fields/Fields.cpp b/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)
+                    };
                 });
         }
     }

diff --git a/src/laser/MultiLaser.H b/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
     };
 }

diff --git a/src/laser/MultiLaser.cpp b/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;
             });