Skip to content

Commit

Permalink
2D/3D agnostic, read base paramters
Browse files Browse the repository at this point in the history
  • Loading branch information
@RevathiJambunathan
RevathiJambunathan committed Aug 1, 2023
1 parent 130aab9 commit 9f7985a
Show file tree
Hide file tree
Showing 4 changed files with 102 additions and 82 deletions.
1 change: 1 addition & 0 deletions Source/Diagnostics/ComputeDiagFunctors/StationFunctor.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ private:
int m_k_index;
int m_max_box_size;
amrex::Box m_buffer_box;
int m_slice_in_domain;
};


Expand Down
163 changes: 83 additions & 80 deletions Source/Diagnostics/ComputeDiagFunctors/StationFunctor.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,85 +30,87 @@ void
StationFunctor::operator ()(amrex::MultiFab& mf_dst, const int dcomp, const int i_buffer) const
{
// 1. First get slice at given z-location
auto& warpx = WarpX::GetInstance();
auto geom = warpx.Geom(m_lev);
// std::unique_ptr< amrex::MultiFab > slice = nullptr;
const int scomp = 0;
const int slice_dir = 2;
bool interpolate = true;

amrex::Vector<int> slice_to_full_ba_map;
std::unique_ptr<amrex::MultiFab> slice = AllocateSlice(slice_to_full_ba_map);

for (amrex::MFIter mfi(*slice, amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
int slice_gid = mfi.index();
int full_gid = slice_to_full_ba_map[slice_gid];
const amrex::Array4<amrex::Real> slice_arr = (*slice)[mfi].array();
const amrex::Array4<amrex::Real const> src_Ex_arr = (*m_arr_mf_src[0])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Ey_arr = (*m_arr_mf_src[1])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Ez_arr = (*m_arr_mf_src[2])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Bx_arr = (*m_arr_mf_src[3])[full_gid].array();
const amrex::Array4<amrex::Real const> src_By_arr = (*m_arr_mf_src[4])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Bz_arr = (*m_arr_mf_src[5])[full_gid].array();

const amrex::Box& tbx = mfi.tilebox();

amrex::ParallelFor( tbx,
[=] AMREX_GPU_DEVICE (int i, int j, int k)
{
slice_arr(i,j,k,0) = src_Ex_arr(i,j,k);
slice_arr(i,j,k,1) = src_Ey_arr(i,j,k);
slice_arr(i,j,k,2) = src_Ez_arr(i,j,k);
slice_arr(i,j,k,3) = src_Bx_arr(i,j,k);
slice_arr(i,j,k,4) = src_By_arr(i,j,k);
slice_arr(i,j,k,5) = src_Bz_arr(i,j,k);
}
);
}

// Define MF with dmap of dst mf, with all 6 components from the slice generated from m_mf_src
const int station_index = static_cast<int> ( ( m_z_location - geom.ProbLo(slice_dir))
/ geom.CellSize(slice_dir) );

amrex::Box slice_box = m_buffer_box;
slice_box.setSmall(slice_dir, station_index);
slice_box.setBig(slice_dir, station_index);

amrex::BoxArray slice_ba(slice_box);
slice_ba.maxSize( m_max_box_size);

std::unique_ptr< amrex::MultiFab > tmp_slice_ptr = nullptr;
const int nghost = 1;
tmp_slice_ptr = std::make_unique< amrex::MultiFab > (slice_ba, mf_dst.DistributionMap(),
slice->nComp(), nghost);
tmp_slice_ptr->setVal(0.);

// Parallel copy slice to tmp_slice MF
const int dcomp_tmp = 0;
amrex::IntVect src_ngrow = amrex::IntVect( AMREX_D_DECL(1,1,1) );
amrex::IntVect dst_ngrow = amrex::IntVect( AMREX_D_DECL(1,1,1) );
ablastr::utils::communication::ParallelCopy(*tmp_slice_ptr, *slice, scomp, dcomp_tmp, slice->nComp(),
slice->nGrowVect(), tmp_slice_ptr->nGrowVect(),
WarpX::do_single_precision_comms);

// MFIter to copy tmp_slice at k-index location in time-based mf
const int k_index = m_k_index;
const int ncomp_dst = mf_dst.nComp();
amrex::MultiFab& tmp = *tmp_slice_ptr;
for (amrex::MFIter mfi(tmp, amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const amrex::Box & tbx = mfi.tilebox();
const amrex::Array4<amrex::Real> src_arr = tmp[mfi].array();
const amrex::Array4<amrex::Real> dst_arr = mf_dst[mfi].array();
amrex::ParallelFor( tbx, ncomp_dst,
[=] AMREX_GPU_DEVICE (int i, int j, int k, int n)
{
dst_arr(i,j,k_index,n) = src_arr(i,j,k,n);
});
if (m_slice_in_domain == 1) {
auto& warpx = WarpX::GetInstance();
auto geom = warpx.Geom(m_lev);
// std::unique_ptr< amrex::MultiFab > slice = nullptr;
const int scomp = 0;
const int slice_dir = 2;
bool interpolate = true;

amrex::Vector<int> slice_to_full_ba_map;
std::unique_ptr<amrex::MultiFab> slice = AllocateSlice(slice_to_full_ba_map);

for (amrex::MFIter mfi(*slice, amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
int slice_gid = mfi.index();
int full_gid = slice_to_full_ba_map[slice_gid];
const amrex::Array4<amrex::Real> slice_arr = (*slice)[mfi].array();
const amrex::Array4<amrex::Real const> src_Ex_arr = (*m_arr_mf_src[0])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Ey_arr = (*m_arr_mf_src[1])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Ez_arr = (*m_arr_mf_src[2])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Bx_arr = (*m_arr_mf_src[3])[full_gid].array();
const amrex::Array4<amrex::Real const> src_By_arr = (*m_arr_mf_src[4])[full_gid].array();
const amrex::Array4<amrex::Real const> src_Bz_arr = (*m_arr_mf_src[5])[full_gid].array();

const amrex::Box& tbx = mfi.tilebox();

amrex::ParallelFor( tbx,
[=] AMREX_GPU_DEVICE (int i, int j, int k)
{
slice_arr(i,j,k,0) = src_Ex_arr(i,j,k);
slice_arr(i,j,k,1) = src_Ey_arr(i,j,k);
slice_arr(i,j,k,2) = src_Ez_arr(i,j,k);
slice_arr(i,j,k,3) = src_Bx_arr(i,j,k);
slice_arr(i,j,k,4) = src_By_arr(i,j,k);
slice_arr(i,j,k,5) = src_Bz_arr(i,j,k);
}
);
}

// Define MF with dmap of dst mf, with all 6 components from the slice generated from m_mf_src
const int station_index = static_cast<int> ( ( m_z_location - geom.ProbLo(slice_dir))
/ geom.CellSize(slice_dir) );

amrex::Box slice_box = m_buffer_box;
slice_box.setSmall(slice_dir, station_index);
slice_box.setBig(slice_dir, station_index);

amrex::BoxArray slice_ba(slice_box);
slice_ba.maxSize( m_max_box_size);

std::unique_ptr< amrex::MultiFab > tmp_slice_ptr = nullptr;
const int nghost = 1;
tmp_slice_ptr = std::make_unique< amrex::MultiFab > (slice_ba, mf_dst.DistributionMap(),
slice->nComp(), nghost);
tmp_slice_ptr->setVal(0.);

// Parallel copy slice to tmp_slice MF
const int dcomp_tmp = 0;
amrex::IntVect src_ngrow = amrex::IntVect( AMREX_D_DECL(1,1,1) );
amrex::IntVect dst_ngrow = amrex::IntVect( AMREX_D_DECL(1,1,1) );
ablastr::utils::communication::ParallelCopy(*tmp_slice_ptr, *slice, scomp, dcomp_tmp, slice->nComp(),
slice->nGrowVect(), tmp_slice_ptr->nGrowVect(),
WarpX::do_single_precision_comms);

// MFIter to copy tmp_slice at k-index location in time-based mf
const int k_index = m_k_index;
const int ncomp_dst = mf_dst.nComp();
amrex::MultiFab& tmp = *tmp_slice_ptr;
for (amrex::MFIter mfi(tmp, amrex::TilingIfNotGPU()); mfi.isValid(); ++mfi)
{
const amrex::Box & tbx = mfi.tilebox();
const amrex::Array4<amrex::Real> src_arr = tmp[mfi].array();
const amrex::Array4<amrex::Real> dst_arr = mf_dst[mfi].array();
amrex::ParallelFor( tbx, ncomp_dst,
[=] AMREX_GPU_DEVICE (int i, int j, int k, int n)
{
dst_arr(i,j,k_index,n) = src_arr(i,j,k,n);
});
}
slice = nullptr;
tmp_slice_ptr = nullptr;
}
slice = nullptr;
tmp_slice_ptr = nullptr;
}

void
Expand All @@ -120,6 +122,7 @@ StationFunctor::PrepareFunctorData (int i_station, bool slice_in_domain, amrex::
m_buffer_box = buffer_box;
m_k_index = k_index;
m_max_box_size = max_box_size;
if (slice_in_domain == true) m_slice_in_domain = 1;
}

std::unique_ptr<amrex::MultiFab>
Expand All @@ -129,8 +132,8 @@ StationFunctor::AllocateSlice (amrex::Vector<int>& slice_to_full_ba_map) const
amrex::Geometry geom = warpx.Geom(m_lev);
// get slice and convert to index space
amrex::RealBox real_slice = geom.ProbDomain();
real_slice.setLo(2, m_z_location);
real_slice.setHi(2, m_z_location);
real_slice.setLo(WARPX_ZINDEX, m_z_location);
real_slice.setHi(WARPX_ZINDEX, m_z_location);
amrex::IntVect slice_lo = amrex::IntVect( AMREX_D_DECL(
static_cast<int>(std::floor((real_slice.lo(0) - geom.ProbLo(0))/geom.CellSize(0))),
static_cast<int>(std::floor((real_slice.lo(1) - geom.ProbLo(1))/geom.CellSize(1))),
Expand Down
1 change: 1 addition & 0 deletions Source/Diagnostics/StationDiagnostics.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ private:
amrex::Box m_buffer_box;
int m_flush_counter = 0;

bool GetZSliceInDomain (const int lev);
};


Expand Down
19 changes: 17 additions & 2 deletions Source/Diagnostics/StationDiagnostics.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ StationDiagnostics::ReadParameters ()
amrex::Abort("StationDiagnostics is not implemented for RZ, yet");
#endif

BaseReadParameters();
const amrex::ParmParse pp_diag_name(m_diag_name);
WARPX_ALWAYS_ASSERT_WITH_MESSAGE(
m_format == "plotfile", "<diag>.format must be plotfile");
Expand Down Expand Up @@ -155,17 +156,31 @@ StationDiagnostics::PrepareFieldDataForOutput ()
for (int lev = 0; lev < nlev; ++lev) {
for (int i = 0; i < num_station_functors; ++i) {
// number of slices = 1
m_all_field_functors[lev][i]->PrepareFunctorData(0, true, m_station_loc, m_buffer_box,
const bool ZSliceInDomain = GetZSliceInDomain(lev);
m_all_field_functors[lev][i]->PrepareFunctorData(0, ZSliceInDomain, m_station_loc, m_buffer_box,
m_slice_counter, m_buffer_size, 0);
}
}
}

bool
StationDiagnostics::GetZSliceInDomain (const int lev)
{
auto & warpx = WarpX::GetInstance();
const amrex::RealBox& prob_domain = warpx.Geom(lev).ProbDomain();
if ( ( m_station_loc <= prob_domain.lo(WARPX_ZINDEX) ) or
( m_station_loc >= prob_domain.hi(WARPX_ZINDEX) ) )
{
return false;
}
return true;
}

void
StationDiagnostics::UpdateBufferData ()
{
m_slice_counter++;
if (GetZSliceInDomain(0))
m_slice_counter++;
}

void
Expand Down

0 comments on commit 9f7985a

Please sign in to comment.