diff options
-rw-r--r-- | Source/Make.WarpX | 4 | ||||
-rw-r--r-- | Source/PhysicalParticleContainer.cpp | 2 | ||||
-rw-r--r-- | Source/WarpXIO.cpp | 4 | ||||
-rw-r--r-- | Source/WarpXParticleContainer.H | 3 | ||||
-rw-r--r-- | Source/WarpXParticleContainer.cpp | 45 |
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 } } |