aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--Source/Make.WarpX4
-rw-r--r--Source/PhysicalParticleContainer.cpp2
-rw-r--r--Source/WarpXIO.cpp4
-rw-r--r--Source/WarpXParticleContainer.H3
-rw-r--r--Source/WarpXParticleContainer.cpp45
5 files changed, 55 insertions, 3 deletions
diff --git a/Source/Make.WarpX b/Source/Make.WarpX
index e5da17e93..5d69dd0fa 100644
--- a/Source/Make.WarpX
+++ b/Source/Make.WarpX
@@ -71,6 +71,10 @@ ifeq ($(USE_PSATD),TRUE)
endif
endif
+ifeq ($(WARPX_RZ),TRUE)
+ DEFINES += -DWARPX_RZ
+endif
+
ifeq ($(STORE_OLD_PARTICLE_ATTRIBS),TRUE)
DEFINES += -DWARPX_STORE_OLD_PARTICLE_ATTRIBS
endif
diff --git a/Source/PhysicalParticleContainer.cpp b/Source/PhysicalParticleContainer.cpp
index 79f132f4d..85644760c 100644
--- a/Source/PhysicalParticleContainer.cpp
+++ b/Source/PhysicalParticleContainer.cpp
@@ -155,7 +155,7 @@ PhysicalParticleContainer::AddGaussianBeam(Real x_m, Real y_m, Real z_m,
std::array<Real, 3> u;
Real weight;
for (long i = 0; i < npart; ++i) {
-#if ( AMREX_SPACEDIM == 3 )
+#if ( AMREX_SPACEDIM == 3 | WARPX_RZ)
weight = q_tot/npart/charge;
Real x = distx(mt);
Real y = disty(mt);
diff --git a/Source/WarpXIO.cpp b/Source/WarpXIO.cpp
index ed2700c75..e68b4a4e8 100644
--- a/Source/WarpXIO.cpp
+++ b/Source/WarpXIO.cpp
@@ -1253,6 +1253,10 @@ WarpX::WritePlotFile () const
particle_varnames.push_back("By");
particle_varnames.push_back("Bz");
+#if (AMREX_SPACEDIM == 2) && WARPX_RZ
+ particle_varnames.push_back("theta");
+#endif
+
#ifdef WARPX_STORE_OLD_PARTICLE_ATTRIBS
particle_varnames.push_back("xold");
particle_varnames.push_back("yold");
diff --git a/Source/WarpXParticleContainer.H b/Source/WarpXParticleContainer.H
index b879f49d9..6b875c13f 100644
--- a/Source/WarpXParticleContainer.H
+++ b/Source/WarpXParticleContainer.H
@@ -11,6 +11,9 @@ struct PIdx
enum { // Particle Attributes stored in amrex::ParticleContainer's struct of array
w = 0, // weight
ux, uy, uz, Ex, Ey, Ez, Bx, By, Bz,
+#if (BL_SPACEDIM == 2) && WARPX_RZ
+ theta, // RZ needs all three position components
+#endif
#ifdef WARPX_STORE_OLD_PARTICLE_ATTRIBS
xold, yold, zold, uxold, uyold, uzold,
#endif
diff --git a/Source/WarpXParticleContainer.cpp b/Source/WarpXParticleContainer.cpp
index c19653287..2bed6e0d3 100644
--- a/Source/WarpXParticleContainer.cpp
+++ b/Source/WarpXParticleContainer.cpp
@@ -21,12 +21,31 @@ void
WarpXParIter::GetPosition (Cuda::DeviceVector<Real>& x, Cuda::DeviceVector<Real>& y, Cuda::DeviceVector<Real>& z) const
{
amrex::ParIter<0,0,PIdx::nattribs>::GetPosition(x, z);
+#ifdef WARPX_RZ
+ const auto& attribs = GetAttribs();
+ const auto& theta = attribs[PIdx::theta];
+ y.resize(x.size());
+ for (unsigned int i=0 ; i < x.size() ; i++) {
+ // The x stored in the particles is actually the radius
+ y[i] = x[i]*std::sin(theta[i]);
+ x[i] = x[i]*std::cos(theta[i]);
+ }
+#else
y.resize(x.size(), std::numeric_limits<Real>::quiet_NaN());
+#endif
}
void
WarpXParIter::SetPosition (const Cuda::DeviceVector<Real>& x, const Cuda::DeviceVector<Real>& y, const Cuda::DeviceVector<Real>& z)
{
+#ifdef WARPX_RZ
+ const auto& attribs = GetAttribs();
+ const auto& theta = attribs[PIdx::theta];
+ for (unsigned int i=0 ; i < x.size() ; i++) {
+ theta[i] = std::atan2(y[i], x[i]);
+ x[i] = std::sqrt(x[i]*x[i] + y[i]*y[i]);
+ }
+#endif
amrex::ParIter<0,0,PIdx::nattribs>::SetPosition(x, z);
}
#endif
@@ -117,6 +136,10 @@ WarpXParticleContainer::AddOneParticle (ParticleTileType& particle_tile,
p.pos(1) = y;
p.pos(2) = z;
#elif (AMREX_SPACEDIM == 2)
+#ifdef WARPX_RZ
+ attribs[PIdx::theta] = std::atan2(y, x);
+ x = std::sqrt(x*x + y*y);
+#endif
p.pos(0) = x;
p.pos(1) = z;
#endif
@@ -157,6 +180,12 @@ WarpXParticleContainer::AddNParticles (int lev,
std::pair<int,int> key {0,0};
auto& particle_tile = GetParticles(lev)[key];
+ std::size_t np = iend-ibegin;
+
+#ifdef WARPX_RZ
+ Vector<Real> theta(np);
+#endif
+
for (int i = ibegin; i < iend; ++i)
{
ParticleType p;
@@ -167,14 +196,17 @@ WarpXParticleContainer::AddNParticles (int lev,
p.pos(1) = y[i];
p.pos(2) = z[i];
#elif (AMREX_SPACEDIM == 2)
+#ifdef WARPX_RZ
+ theta[i-ibegin] = std::atan2(y[i], x[i]);
+ p.pos(0) = std::sqrt(x[i]*x[i] + y[i]*y[i]);
+#else
p.pos(0) = x[i];
+#endif
p.pos(1) = z[i];
#endif
particle_tile.push_back(p);
}
- std::size_t np = iend-ibegin;
-
if (np > 0)
{
particle_tile.push_back_real(PIdx::w , weight + ibegin, weight + iend);
@@ -184,7 +216,16 @@ WarpXParticleContainer::AddNParticles (int lev,
for (int comp = PIdx::uz+1; comp < PIdx::nattribs; ++comp)
{
+#ifdef WARPX_RZ
+ if (comp == PIdx::theta) {
+ particle_tile.push_back_real(comp, theta.data, theta.data+np);
+ }
+ else {
+ particle_tile.push_back_real(comp, np, 0.0);
+ }
+#else
particle_tile.push_back_real(comp, np, 0.0);
+#endif
}
}