Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -35,10 +35,8 @@ def __init__(self):
self.spacecraft_capacitance = None

# shortcuts
self.dir_r, self.dir_z = (
libwarpx.libwarpx_so.Direction(0),
libwarpx.libwarpx_so.Direction(2),
)
Direction = libwarpx.libwarpx_so.Direction
self.dir_r, self.dir_z = (Direction.r, Direction.z)

def correct_space_charge_fields(self, q=None):
"""
Expand Down
22 changes: 11 additions & 11 deletions Source/Diagnostics/BTDiagnostics.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -615,7 +615,7 @@ BTDiagnostics::UpdateVarnamesForRZopenPMD ()
auto & warpx = WarpX::GetInstance();
auto & fields = warpx.m_fields;
using ablastr::fields::Direction;
const int ncomp_multimodefab = fields.get(FieldType::Efield_aux, Direction{0}, 0)->nComp();
const int ncomp_multimodefab = fields.get(FieldType::Efield_aux, Direction::r, 0)->nComp();
const int ncomp = ncomp_multimodefab;


Expand Down Expand Up @@ -677,7 +677,7 @@ BTDiagnostics::InitializeFieldFunctorsRZopenPMD (int lev)

auto & warpx = WarpX::GetInstance();
auto & fields = warpx.m_fields;
const int ncomp_multimodefab = fields.get(FieldType::Efield_aux, Direction{0}, 0)->nComp();
const int ncomp_multimodefab = fields.get(FieldType::Efield_aux, Direction::r, 0)->nComp();
const int ncomp = ncomp_multimodefab;
// Clear any pre-existing vector to release stored data
// This ensures that when domain is load-balanced, the functors point
Expand All @@ -703,23 +703,23 @@ BTDiagnostics::InitializeFieldFunctorsRZopenPMD (int lev)
const auto m_cell_center_functors_at_lev_size = static_cast<int>(m_cell_center_functors.at(lev).size());
for (int comp=0; comp<m_cell_center_functors_at_lev_size; comp++){
if ( m_cellcenter_varnames_fields[comp] == "Er" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Efield_aux, Direction{0}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Efield_aux, Direction::r, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "Et" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Efield_aux, Direction{1}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Efield_aux, Direction::theta, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "Ez" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Efield_aux, Direction{2}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Efield_aux, Direction::z, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "Br" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Bfield_aux, Direction{0}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Bfield_aux, Direction::r, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "Bt" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Bfield_aux, Direction{1}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Bfield_aux, Direction::theta, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "Bz" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Bfield_aux, Direction{2}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::Bfield_aux, Direction::z, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "jr" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::current_fp, Direction{0}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::current_fp, Direction::r, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "jt" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::current_fp, Direction{1}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::current_fp, Direction::theta, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "jz" ){
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::current_fp, Direction{2}, lev), lev, m_crse_ratio, false, ncomp);
m_cell_center_functors[lev][comp] = std::make_unique<CellCenterFunctor>(fields.get(FieldType::current_fp, Direction::z, lev), lev, m_crse_ratio, false, ncomp);
} else if ( m_cellcenter_varnames_fields[comp] == "rho" ){
m_cell_center_functors[lev][comp] = std::make_unique<RhoFunctor>(lev, m_crse_ratio, false, -1, false, ncomp);
}
Expand Down
13 changes: 5 additions & 8 deletions Source/Diagnostics/FullDiagnostics.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -370,7 +370,7 @@ FullDiagnostics::InitializeFieldFunctorsRZopenPMD (int lev)
using ablastr::fields::Direction;

auto & warpx = WarpX::GetInstance();
const int ncomp_multimodefab = warpx.m_fields.get(FieldType::Efield_aux, Direction{0}, 0)->nComp();
const int ncomp_multimodefab = warpx.m_fields.get(FieldType::Efield_aux, Direction::r, 0)->nComp();
// Make sure all multifabs have the same number of components
for (int dim=0; dim<3; dim++){
AMREX_ALWAYS_ASSERT(
Expand Down Expand Up @@ -573,7 +573,7 @@ FullDiagnostics::AddRZModesToDiags (int lev)
if (!m_dump_rz_modes) { return; }

auto & warpx = WarpX::GetInstance();
const int ncomp_multimodefab = warpx.m_fields.get(FieldType::Efield_aux, Direction{0}, 0)->nComp();
const int ncomp_multimodefab = warpx.m_fields.get(FieldType::Efield_aux, Direction::r, 0)->nComp();
// Make sure all multifabs have the same number of components
for (int dim=0; dim<3; dim++){
AMREX_ALWAYS_ASSERT(
Expand Down Expand Up @@ -616,26 +616,23 @@ FullDiagnostics::AddRZModesToDiags (int lev)
m_all_field_functors[lev].push_back(std::make_unique<CellCenterFunctor>(
warpx.m_fields.get(FieldType::Efield_aux, Direction{dim}, lev), lev,
m_crse_ratio, false, ncomp_multimodefab));
AddRZModesToOutputNames(std::string("E") + coord[dim],
warpx.m_fields.get(FieldType::Efield_aux, Direction{0}, 0)->nComp());
AddRZModesToOutputNames(std::string("E") + coord[dim], ncomp_multimodefab);
}
// B
for (int dim=0; dim<3; dim++){
// 3 components, r theta z
m_all_field_functors[lev].push_back(std::make_unique<CellCenterFunctor>(
warpx.m_fields.get(FieldType::Bfield_aux, Direction{dim}, lev), lev,
m_crse_ratio, false, ncomp_multimodefab));
AddRZModesToOutputNames(std::string("B") + coord[dim],
warpx.m_fields.get(FieldType::Bfield_aux, Direction{0}, 0)->nComp());
AddRZModesToOutputNames(std::string("B") + coord[dim], ncomp_multimodefab);
}
// j
for (int dim=0; dim<3; dim++){
// 3 components, r theta z
m_all_field_functors[lev].push_back(std::make_unique<JFunctor>(
dim, lev, m_crse_ratio, false, deposit_current, ncomp_multimodefab));
deposit_current = false;
AddRZModesToOutputNames(std::string("J") + coord[dim],
warpx.m_fields.get(FieldType::current_fp,Direction{0},0)->nComp());
AddRZModesToOutputNames(std::string("J") + coord[dim], ncomp_multimodefab);
}
// divE
if (divE_requested) {
Expand Down
39 changes: 37 additions & 2 deletions Source/Python/MultiFabRegister.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,43 @@ void init_MultiFabRegister (py::module & m)
{
using namespace ablastr::fields;

py::class_<ablastr::fields::Direction>(m, "Direction")
.def(py::init<int>());
py::class_<ablastr::fields::Direction> pyDirection(m, "Direction");
pyDirection
.def(py::init<int>())
.def(py::init<std::string>())
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
.def_property_readonly_static("r", [](py::object /* self */) {
return ablastr::fields::Direction::r;
})
#endif
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
.def_property_readonly_static("theta", [](py::object /* self */) {
return ablastr::fields::Direction::theta;
})
#endif

#if defined(WARPX_DIM_RSPHERE)
.def_property_readonly_static("phi", [](py::object /* self */) {
return ablastr::fields::Direction::phi;
})
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
.def_property_readonly_static("x", [](py::object /* self */) {
return ablastr::fields::Direction::x;
})
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
.def_property_readonly_static("y", [](py::object /* self */) {
return ablastr::fields::Direction::y;
})
#endif
#if !defined(WARPX_DIM_RSPHERE)
.def_property_readonly_static("z", [](py::object /* self */) {
return ablastr::fields::Direction::z;
})
#endif
;
py::implicitly_convertible<std::string, ablastr::fields::Direction>();

py::class_<ablastr::fields::MultiFabRegister>(m, "MultiFabRegister")

Expand Down
121 changes: 118 additions & 3 deletions Source/ablastr/fields/MultiFabRegister.H
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
#include <map>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <vector>
Expand Down Expand Up @@ -58,21 +59,120 @@ namespace ablastr::fields
* Because of different staggering, the components of vector/tensor fields are stored
* in separate (i)MultiFab.
*
* Cartesian: always x,y,z in any ND
* RZ: r,t(heta),z (r,z over azimuthal modes)
* RCylinder: r,t(heta),z
* RSphere: r,t(heta),p(hi)
* @see https://warpx.readthedocs.io/en/latest/developers/dimensionality.html
*
* @todo: synchronize with AMReX "enum class Direction"
*/
struct Direction
class Direction
{
int dir = 0;

public:
constexpr explicit Direction (int d) : dir(d) {}

#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
static const Direction r;
static const Direction theta;
#endif
#if defined(WARPX_DIM_RSPHERE)
static const Direction phi;
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
static const Direction x;
static const Direction y;
#endif
#if !defined(WARPX_DIM_RSPHERE)
static const Direction z;
#endif

bool operator<(const Direction& other) const
{
return other.dir < this->dir;
}

operator std::string() const
{
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
if (dir == r) { return "r"; }
if (dir == theta) { return "theta"; }
#endif
#if defined(WARPX_DIM_RSPHERE)
if (dir == phi) { return "phi"; }
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
if (dir == x) { return "x"; }
if (dir == y) { return "y"; }
#endif
#if !defined(WARPX_DIM_RSPHERE)
if (dir == z) { return "z"; }
#endif
throw std::runtime_error("invalid direction: " + std::to_string(dir));
return std::to_string(dir);
}

Direction (std::string const & s)
{
dir = -1;
#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
if (s.compare(r) == 0) { dir = r; }
if (s.compare(theta) == 0) { dir = theta; }
#endif
#if defined(WARPX_DIM_RSPHERE)
if (s.compare(phi) == 0) { dir = phi; }
#endif
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
if (s.compare(x) == 0) { dir = x; }
if (s.compare(y) == 0) { dir = y; }
#endif
#if !defined(WARPX_DIM_RSPHERE)
if (s.compare(z) == 0) { dir = z; }
#endif

if (dir == -1) {
throw std::runtime_error("invalid direction: " + s);
}
}

// e.g., "x"
Direction (char const * c) : dir{Direction{c}} {}

// e.g., 'x'
Direction (char const c) : dir{Direction{c}} {}

// rule of five: define default constructors and the destructor
// because we defined a special one above for char
~Direction () = default;
Direction (const Direction&) = default;
Direction& operator= (const Direction&) = default;
Direction (Direction&&) = default;
Direction& operator= (Direction&&) = default;

/* TODO: just temporary int compatibility */
operator int() const { return dir; }

};

#if defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER) || defined(WARPX_DIM_RSPHERE)
inline constexpr Direction Direction::r = Direction{0};
inline constexpr Direction Direction::theta = Direction{1};
#endif
#if defined(WARPX_DIM_RSPHERE)
inline constexpr Direction Direction::phi = Direction{2};
#endif

#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
inline constexpr Direction Direction::x = Direction{0};
inline constexpr Direction Direction::y = Direction{1};
#endif

#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z) || defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER)
inline constexpr Direction Direction::z = Direction{2};
#endif

/** A scalar field (a MultiFab)
*
* Note: might still have components, e.g., for copies at different times.
Expand Down Expand Up @@ -822,8 +922,23 @@ namespace ablastr::fields
MultiFabOwner
> m_mf_register;

/** the three directions of a vector field */
std::vector<Direction> m_all_dirs = {Direction{0}, Direction{1}, Direction{2}};
public:
/** The directions of a vector field as stored in the simulation.
*
* Cartesian: always x,y,z in any ND
* RZ: r,t(heta),z (r,z over azimuthal modes)
* RCylinder: r,t(heta),z
* RSphere: r,t(heta),p(hi)
* @see https://warpx.readthedocs.io/en/latest/developers/dimensionality.html
*/
static inline std::vector<Direction> m_all_dirs =
#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_XZ) || defined(WARPX_DIM_1D_Z)
{Direction::x, Direction::y, Direction::z};
#elif defined(WARPX_DIM_RZ) || defined(WARPX_DIM_RCYLINDER)
{Direction::r, Direction::theta, Direction::z};
#elif defined(WARPX_DIM_RSPHERE)
{Direction::r, Direction::theta, Direction::phi};
#endif
};

/** Little temporary helper function to pass temporary MultiFabs as VectorField.
Expand Down
14 changes: 4 additions & 10 deletions Source/ablastr/fields/MultiFabRegister.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -628,15 +628,11 @@ namespace ablastr::fields
int level
) const
{
// Add the suffix for the direction [x] or [y] or [z]
// note: since Cartesian is not correct for all our supported geometries,
// in the future we might want to break this to "[dir=0/1/2]".
// This will be a breaking change for (Python) users that rely on that string.
constexpr int x_in_ascii = 120;
std::string const component_name{char(x_in_ascii + dir.dir)};
// Add the suffix for the direction [dir=x] or [dir=y] or [dir=z] or [dir=r]
std::string const component_name = dir;
return mf_name(
name
.append("[")
.append("[dir=")
Copy link
Member Author

@ax3l ax3l Sep 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@roelof-groenewald @RemiLehe @dpgrote just highlighting that this is a breaking change if someone uses the "raw old strings" still in Python instead of the APIs we defined a year ago. I do not think someone does and our wrappers are mostly used anyway, which do the right thing...

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@n01r I think that you may have some scripts that use this. Just flagging the breaking change here.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should the use of the internal names be deprecated?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, ideally. This is mostly to make sure that list() looks understandable.

Do we use internal names still somewhere?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not used anywhere within WarpX, but the Python interface is still there

        .def("multifab",
             [](WarpX & wx, std::string internal_name) {

There is a note that this is deprecated, but this should just be removed.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. Let me remove this in a follow-up PR, since this one is all green.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Posted in #6180

.append(component_name)
.append("]"),
level
Expand All @@ -648,12 +644,10 @@ namespace ablastr::fields
std::array< std::unique_ptr<amrex::MultiFab>, 3 > const & old_vectorfield
)
{
std::vector<Direction> const all_dirs = {Direction{0}, Direction{1}, Direction{2}};

VectorField field_on_level;

// insert components
for (auto const dir : {0, 1, 2})
for (auto const dir : MultiFabRegister::m_all_dirs)
{
field_on_level[Direction{dir}] = old_vectorfield[dir].get();
}
Expand Down
Loading