2 files changed, 35 insertions, 2 deletions
diff --git a/Source/Particles/Sorting/Partition.cpp b/Source/Particles/Sorting/Partition.cpp
index 49d4d25a6..6874e2df6 100644
--- a/Source/Particles/Sorting/Partition.cpp
+++ b/Source/Particles/Sorting/Partition.cpp
@@ -54,9 +54,20 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     auto& aos = pti.GetArrayOfStructs();
     ParticleType * AMREX_RESTRICT pstructs = &(aos[0]);
     int * inexflag_ptr = inexflag.dataPtr();
-    amrex::ParallelFor( pti.numParticles(),
+    const Geometry& geom = Geom(lev);
+    const Real prob_lo_x = geom.ProbLo(0);
+    const Real prob_lo_y = geom.ProbLo(1);
+    const Real prob_lo_z = geom.ProbLo(2);
+    const Real inv_dx = geom.InvCellSize(0);
+    const Real inv_dy = geom.InvCellSize(1);
+    const Real inv_dz = geom.InvCellSize(2);
+    const IntVect domain_small_end = geom.Domain().smallEnd();
+    amrex::ParallelFor( np,
          [=] AMREX_GPU_DEVICE (long i) {
-             const IntVect& iv = Index(pstructs[i], lev);
+
+             const IntVect& iv = findParticleIndex(pstructs[i],
+                 prob_lo_x, prob_lo_y, prob_lo_z,
+                 inv_dx, inv_dy, inv_dz, domain_small_end );
              inexflag_ptr[i] = msk(iv);
          }
      );
diff --git a/Source/Particles/Sorting/SortingUtils.H b/Source/Particles/Sorting/SortingUtils.H
index d072ebd2a..7dec21446 100644
--- a/Source/Particles/Sorting/SortingUtils.H
+++ b/Source/Particles/Sorting/SortingUtils.H
@@ -15,4 +15,26 @@ void fillWithConsecutiveIntegers( amrex::Gpu::ManagedDeviceVector<long>& v ) {
 #endif
 }
 
+// TODO: Add documentation
+template< typename PType >
+AMREX_GPU_HOST_DEVICE AMREX_INLINE
+amrex::IntVect findParticleIndex ( PType p,
+                    const amrex::Real prob_lo_x,
+                    const amrex::Real prob_lo_y,
+                    const amrex::Real prob_lo_z,
+                    const amrex::Real inv_dx,
+                    const amrex::Real inv_dy,
+                    const amrex::Real inv_dz,
+                    const amrex::IntVect domain_small_end )
+{
+    amrex::IntVect iv;
+    AMREX_D_TERM(iv[0]=static_cast<int>(floor((p.m_rdata.pos[0]-prob_lo_x)*inv_dx));,
+                 iv[1]=static_cast<int>(floor((p.m_rdata.pos[1]-prob_lo_y)*inv_dy));,
+                 iv[2]=static_cast<int>(floor((p.m_rdata.pos[2]-prob_lo_z)*inv_dz)););
+
+    iv += domain_small_end;
+
+    return iv;
+};
+
 #endif // WARPX_PARTICLES_SORTING_SORTINGUTILS_H_