Skip to content

Commit

Permalink
major refactor of magnetostatic solver
Browse files Browse the repository at this point in the history
  • Loading branch information
roelof-groenewald committed Aug 27, 2024
1 parent 874a481 commit 6d04381
Show file tree
Hide file tree
Showing 10 changed files with 360 additions and 211 deletions.
6 changes: 3 additions & 3 deletions Python/pywarpx/fields.py
Original file line number Diff line number Diff line change
Expand Up @@ -701,23 +701,23 @@ def GFPWrapper(level=0, include_ghosts=False):

def AxFPWrapper(level=0, include_ghosts=False):
return _MultiFABWrapper(
mf_name="vector_potential_fp_nodal[x]",
mf_name="Afield_fp_nodal[x]",
level=level,
include_ghosts=include_ghosts,
)


def AyFPWrapper(level=0, include_ghosts=False):
return _MultiFABWrapper(
mf_name="vector_potential_fp_nodal[y]",
mf_name="Afield_fp_nodal[y]",
level=level,
include_ghosts=include_ghosts,
)


def AzFPWrapper(level=0, include_ghosts=False):
return _MultiFABWrapper(
mf_name="vector_potential_fp_nodal[z]",
mf_name="Afield_fp_nodal[z]",
level=level,
include_ghosts=include_ghosts,
)
Expand Down
188 changes: 188 additions & 0 deletions Source/FieldSolver/FiniteDifferenceSolver/EvolveB.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -484,3 +484,191 @@ void FiniteDifferenceSolver::EvolveBCylindrical (
}

#endif // corresponds to ifndef WARPX_DIM_RZ


void FiniteDifferenceSolver::CalcBfromVectorPotential (
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>& Bfield,
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3> const& Afield,
std::array<std::unique_ptr<amrex::MultiFab>, 3> const& edge_lengths,
const int lev )
{
if (m_fdtd_algo != ElectromagneticSolverAlgo::CKC && m_fdtd_algo != ElectromagneticSolverAlgo::PSATD) {
#ifdef WARPX_DIM_RZ
CalcBfromVectorPotentialCylindrical <CylindricalYeeAlgorithm> (
Bfield, Afield, edge_lengths, lev
);
#else
CalcBfromVectorPotentialCartesian <CartesianYeeAlgorithm> (
Bfield, Afield, edge_lengths, lev
);
#endif
} else {
WARPX_ABORT_WITH_MESSAGE("CalcBfromVectorPotential: Unknown algorithm");
}
}

#ifdef WARPX_DIM_RZ
template<typename T_Algo>
void FiniteDifferenceSolver::CalcBfromVectorPotentialCylindrical (
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>& Bfield,
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3> const& Afield,
[[maybe_unused]] std::array< std::unique_ptr<amrex::MultiFab>, 3 > const& edge_lengths,
[[maybe_unused]] const int lev
)
{
// Loop through the grids, and over the tiles within each grid
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for ( MFIter mfi(*Bfield[0], TilingIfNotGPU()); mfi.isValid(); ++mfi ) {

// Extract field data for this grid/tile
Array4<Real> const& Br = Bfield[0]->array(mfi);
Array4<Real> const& Bt = Bfield[1]->array(mfi);
Array4<Real> const& Bz = Bfield[2]->array(mfi);
Array4<Real> const& Ar = Afield[0]->array(mfi);
Array4<Real> const& At = Afield[1]->array(mfi);
Array4<Real> const& Az = Afield[2]->array(mfi);

// Extract stencil coefficients
Real const * const AMREX_RESTRICT coefs_r = m_stencil_coefs_r.dataPtr();
auto const n_coefs_r = static_cast<int>(m_stencil_coefs_r.size());
Real const * const AMREX_RESTRICT coefs_z = m_stencil_coefs_z.dataPtr();
auto const n_coefs_z = static_cast<int>(m_stencil_coefs_z.size());

// Extract cylindrical specific parameters
Real const dr = m_dr;
int const nmodes = m_nmodes;
Real const rmin = m_rmin;

// Extract tileboxes for which to loop
Box const& tbr = mfi.tilebox(Bfield[0]->ixType().toIntVect());
Box const& tbt = mfi.tilebox(Bfield[1]->ixType().toIntVect());
Box const& tbz = mfi.tilebox(Bfield[2]->ixType().toIntVect());

// Loop over the cells and update the fields
amrex::ParallelFor(tbr, tbt, tbz,

[=] AMREX_GPU_DEVICE (int i, int j, int /*k*/){
Real const r = rmin + i*dr; // r on nodal point (Br is nodal in r)
if (r != 0) { // Off-axis, regular Maxwell equations
Br(i, j, 0, 0) += T_Algo::UpwardDz(At, coefs_z, n_coefs_z, i, j, 0, 0); // Mode m=0
for (int m=1; m<nmodes; m++) { // Higher-order modes
Br(i, j, 0, 2*m-1) += (
T_Algo::UpwardDz(At, coefs_z, n_coefs_z, i, j, 0, 2*m-1)
- m * Az(i, j, 0, 2*m )/r ); // Real part
Br(i, j, 0, 2*m ) += (
T_Algo::UpwardDz(At, coefs_z, n_coefs_z, i, j, 0, 2*m )
+ m * Az(i, j, 0, 2*m-1)/r ); // Imaginary part
}
} else { // r==0: On-axis corrections
// Ensure that Br remains 0 on axis (except for m=1)
Br(i, j, 0, 0) = 0.; // Mode m=0
for (int m=1; m<nmodes; m++) { // Higher-order modes
if (m == 1){
// For m==1, Az is linear in r, for small r
// Therefore, the formula below regularizes the singularity
Br(i, j, 0, 2*m-1) += (
T_Algo::UpwardDz(At, coefs_z, n_coefs_z, i, j, 0, 2*m-1)
- m * Az(i+1, j, 0, 2*m )/dr ); // Real part
Br(i, j, 0, 2*m ) += (
T_Algo::UpwardDz(At, coefs_z, n_coefs_z, i, j, 0, 2*m )
+ m * Az(i+1, j, 0, 2*m-1)/dr ); // Imaginary part
} else {
Br(i, j, 0, 2*m-1) = 0.;
Br(i, j, 0, 2*m ) = 0.;
}
}
}
},

[=] AMREX_GPU_DEVICE (int i, int j, int /*k*/){
Bt(i, j, 0, 0) += (
T_Algo::UpwardDr(Az, coefs_r, n_coefs_r, i, j, 0, 0)
- T_Algo::UpwardDz(Ar, coefs_z, n_coefs_z, i, j, 0, 0)); // Mode m=0
for (int m=1 ; m<nmodes ; m++) { // Higher-order modes
Bt(i, j, 0, 2*m-1) += (
T_Algo::UpwardDr(Az, coefs_r, n_coefs_r, i, j, 0, 2*m-1)
- T_Algo::UpwardDz(Ar, coefs_z, n_coefs_z, i, j, 0, 2*m-1)); // Real part
Bt(i, j, 0, 2*m ) += (
T_Algo::UpwardDr(Az, coefs_r, n_coefs_r, i, j, 0, 2*m )
- T_Algo::UpwardDz(Ar, coefs_z, n_coefs_z, i, j, 0, 2*m )); // Imaginary part
}
},

[=] AMREX_GPU_DEVICE (int i, int j, int /*k*/){
Real const r = rmin + (i + 0.5_rt)*dr; // r on a cell-centered grid (Bz is cell-centered in r)
Bz(i, j, 0, 0) += ( - T_Algo::UpwardDrr_over_r(At, r, dr, coefs_r, n_coefs_r, i, j, 0, 0));
for (int m=1 ; m<nmodes ; m++) { // Higher-order modes
Bz(i, j, 0, 2*m-1) += ( m * Ar(i, j, 0, 2*m )/r
- T_Algo::UpwardDrr_over_r(At, r, dr, coefs_r, n_coefs_r, i, j, 0, 2*m-1)); // Real part
Bz(i, j, 0, 2*m ) += (-m * Ar(i, j, 0, 2*m-1)/r
- T_Algo::UpwardDrr_over_r(At, r, dr, coefs_r, n_coefs_r, i, j, 0, 2*m )); // Imaginary part
}
}

);
}
}
#else
template<typename T_Algo>
void FiniteDifferenceSolver::CalcBfromVectorPotentialCartesian (
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>& Bfield,
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3> const& Afield,
[[maybe_unused]] std::array< std::unique_ptr<amrex::MultiFab>, 3 > const& edge_lengths,
[[maybe_unused]] const int lev
)
{
// Loop through the grids, and over the tiles within each grid
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
for ( MFIter mfi(*Bfield[0], TilingIfNotGPU()); mfi.isValid(); ++mfi ) {
// Extract field data for this grid/tile
Array4<Real> const& Bx = Bfield[0]->array(mfi);
Array4<Real> const& By = Bfield[1]->array(mfi);
Array4<Real> const& Bz = Bfield[2]->array(mfi);
Array4<Real> const& Ax = Afield[0]->array(mfi);
Array4<Real> const& Ay = Afield[1]->array(mfi);
Array4<Real> const& Az = Afield[2]->array(mfi);

// Extract stencil coefficients
Real const * const AMREX_RESTRICT coefs_x = m_stencil_coefs_x.dataPtr();
auto const n_coefs_x = static_cast<int>(m_stencil_coefs_x.size());
Real const * const AMREX_RESTRICT coefs_y = m_stencil_coefs_y.dataPtr();
auto const n_coefs_y = static_cast<int>(m_stencil_coefs_y.size());
Real const * const AMREX_RESTRICT coefs_z = m_stencil_coefs_z.dataPtr();
auto const n_coefs_z = static_cast<int>(m_stencil_coefs_z.size());

// Extract tileboxes for which to loop
Box const& tbx = mfi.tilebox(Bfield[0]->ixType().toIntVect());
Box const& tby = mfi.tilebox(Bfield[1]->ixType().toIntVect());
Box const& tbz = mfi.tilebox(Bfield[2]->ixType().toIntVect());

// Loop over the cells and update the fields
amrex::ParallelFor(tbx, tby, tbz,

[=] AMREX_GPU_DEVICE (int i, int j, int k){

Bx(i, j, k) += T_Algo::UpwardDz(Ay, coefs_z, n_coefs_z, i, j, k)
- T_Algo::UpwardDy(Az, coefs_y, n_coefs_y, i, j, k);

},

[=] AMREX_GPU_DEVICE (int i, int j, int k){

By(i, j, k) += T_Algo::UpwardDx(Az, coefs_x, n_coefs_x, i, j, k)
- T_Algo::UpwardDz(Ax, coefs_z, n_coefs_z, i, j, k);

},

[=] AMREX_GPU_DEVICE (int i, int j, int k){

Bz(i, j, k) += T_Algo::UpwardDy(Ax, coefs_y, n_coefs_y, i, j, k)
- T_Algo::UpwardDx(Ay, coefs_x, n_coefs_x, i, j, k);

}
);
}
}
#endif // corresponds to ifndef WARPX_DIM_RZ
32 changes: 32 additions & 0 deletions Source/FieldSolver/FiniteDifferenceSolver/FiniteDifferenceSolver.H
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,22 @@ class FiniteDifferenceSolver
std::array< std::unique_ptr<amrex::MultiFab>, 3 > const& edge_lengths,
int lev );

/**
* \brief Calculation of the magnetic field given a vector potential
* field from B = ∇ x A.
*
* \param[out] Bfield vector of B-field MultiFabs at a given level
* \param[in] Afield vector of vector potential MultiFabs at a given level
* \param[in] edge_lengths length of edges along embedded boundaries
* \param[in] lev level number for the calculation
*/
void CalcBfromVectorPotential (
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>& Bfield,
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3> const& Afield,
std::array< std::unique_ptr<amrex::MultiFab>, 3 > const& edge_lengths,
int lev
);

private:

int m_fdtd_algo;
Expand Down Expand Up @@ -255,6 +271,14 @@ class FiniteDifferenceSolver
int lev
);

template<typename T_Algo>
void CalcBfromVectorPotentialCylindrical (
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>& Bfield,
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3> const& Afield,
std::array< std::unique_ptr<amrex::MultiFab>, 3 > const& edge_lengths,
int lev
);

#else
template< typename T_Algo >
void EvolveBCartesian (
Expand Down Expand Up @@ -359,6 +383,14 @@ class FiniteDifferenceSolver
std::array< std::unique_ptr<amrex::MultiFab>, 3 > const& edge_lengths,
int lev
);

template<typename T_Algo>
void CalcBfromVectorPotentialCartesian (
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>& Bfield,
amrex::Array<std::unique_ptr<amrex::MultiFab>, 3> const& Afield,
std::array< std::unique_ptr<amrex::MultiFab>, 3 > const& edge_lengths,
int lev
);
#endif

};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ FiniteDifferenceSolver::FiniteDifferenceSolver (
m_grid_type{grid_type}
{
// return if not FDTD
if (fdtd_algo == ElectromagneticSolverAlgo::None || fdtd_algo == ElectromagneticSolverAlgo::PSATD) {
if (fdtd_algo == ElectromagneticSolverAlgo::PSATD) {
return;
}

Expand All @@ -45,7 +45,8 @@ FiniteDifferenceSolver::FiniteDifferenceSolver (
m_dr = cell_size[0];
m_nmodes = WarpX::n_rz_azimuthal_modes;
m_rmin = WarpX::GetInstance().Geom(0).ProbLo(0);
if (fdtd_algo == ElectromagneticSolverAlgo::Yee ||
if (fdtd_algo == ElectromagneticSolverAlgo::None ||
fdtd_algo == ElectromagneticSolverAlgo::Yee ||
fdtd_algo == ElectromagneticSolverAlgo::HybridPIC ) {
CylindricalYeeAlgorithm::InitializeStencilCoefficients( cell_size,
m_h_stencil_coefs_r, m_h_stencil_coefs_z );
Expand All @@ -68,7 +69,8 @@ FiniteDifferenceSolver::FiniteDifferenceSolver (
CartesianNodalAlgorithm::InitializeStencilCoefficients( cell_size,
m_h_stencil_coefs_x, m_h_stencil_coefs_y, m_h_stencil_coefs_z );

} else if (fdtd_algo == ElectromagneticSolverAlgo::Yee ||
} else if (fdtd_algo == ElectromagneticSolverAlgo::None ||
fdtd_algo == ElectromagneticSolverAlgo::Yee ||
fdtd_algo == ElectromagneticSolverAlgo::ECT ||
fdtd_algo == ElectromagneticSolverAlgo::HybridPIC) {

Expand Down
30 changes: 4 additions & 26 deletions Source/FieldSolver/MagnetostaticSolver/MagnetostaticSolver.H
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
/* Copyright 2022 S. Eric Clark, LLNL
/* Copyright 2022 The WarpX Community
*
* This file is part of WarpX.
*
* Authors: S. Eric Clark (LLNL)
*
* License: BSD-3-Clause-LBNL
*/
#ifndef WARPX_MAGNETOSTATICSOLVER_H_
Expand All @@ -12,6 +14,7 @@
#include <AMReX_MLMG.H>
#include <AMReX_REAL.H>


namespace MagnetostaticSolver {

/** Boundary Handler for the Vector Potential Poisson Solver
Expand All @@ -27,31 +30,6 @@ namespace MagnetostaticSolver {

void defineVectorPotentialBCs ();
};

/** use amrex to directly calculate the magnetic field since with EB's the
*
* simple finite difference scheme in WarpX::computeE sometimes fails
*/
class EBCalcBfromVectorPotentialPerLevel {
private:
const amrex::Vector<amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>>& m_b_field;
const amrex::Vector<amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>>& m_grad_buf_e_stag;
const amrex::Vector<amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>>& m_grad_buf_b_stag;

public:
EBCalcBfromVectorPotentialPerLevel(const amrex::Vector<amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>>& b_field,
const amrex::Vector<amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>>& grad_buf_e_stag,
const amrex::Vector<amrex::Array<std::unique_ptr<amrex::MultiFab>, 3>>& grad_buf_b_stag)
: m_b_field(b_field),
m_grad_buf_e_stag(grad_buf_e_stag),
m_grad_buf_b_stag(grad_buf_b_stag)
{}

void operator()(amrex::Array<std::unique_ptr<amrex::MLMG>,3> & mlmg, int lev);

// Function to perform interpolation from cell edges to cell faces
void doInterp(const std::unique_ptr<amrex::MultiFab> &src, const std::unique_ptr<amrex::MultiFab> &dst);
};
} // namespace MagnetostaticSolver

#endif //WARPX_MAGNETOSTATICSOLVER_H_
Loading

0 comments on commit 6d04381

Please sign in to comment.