diff --git a/runtime/domain/block_forest.cpp b/runtime/domain/block_forest.cpp
index 8da11deaa0568af1a9b56c5e476c7072c7456825..9092cbe8b42773e0e3ba8352c3576711626f8221 100644
--- a/runtime/domain/block_forest.cpp
+++ b/runtime/domain/block_forest.cpp
@@ -22,11 +22,12 @@ namespace pairs {
 
 void BlockForest::updateNeighborhood() {
     auto me = mpi::MPIManager::instance()->rank();
+    this->nranks = 0;
+    this->total_aabbs = 0;
+
     ranks.clear();
     naabbs.clear();
     aabbs.clear();
-    this->nranks = 0;
-    this->total_aabbs = 0;
 
     for(auto& iblock: *forest) {
         auto block = static_cast<blockforest::Block *>(&iblock);
@@ -80,7 +81,6 @@ void BlockForest::copyRuntimeArray(const std::string& name, void *dest, const in
     void *src = name.compare('ranks') ? ranks.data() :
                 name.compare('naabbs') ? vec_naabbs.data() :
                 name.compare('rank_offsets') ? offsets :
-                name.compare('pbc') ? vec_pbc.data() :
                 name.compare('aabbs') ? vec_aabbs.data() :
                 name.compare('subdom') ? subdom;
 
diff --git a/src/pairs/sim/domain_partitioning.py b/src/pairs/sim/domain_partitioning.py
index 1c88c4c4f750457d79ea5df7ee0fccdfe8317ab5..6adaa85c48b2c9562a5165db97195d578474d046 100644
--- a/src/pairs/sim/domain_partitioning.py
+++ b/src/pairs/sim/domain_partitioning.py
@@ -83,7 +83,6 @@ class BlockForest:
         self.ranks              = sim.add_static_array('ranks', [self.nranks_capacity], Types.Int32)
         self.naabbs             = sim.add_static_array('naabbs', [self.nranks_capacity], Types.Int32)
         self.offsets            = sim.add_static_array('rank_offsets', [self.nranks_capacity], Types.Int32)
-        self.pbc                = sim.add_static_array('pbc', [self.aabb_capacity, 3], Types.Int32)
         self.aabbs              = sim.add_static_array('aabbs', [self.aabb_capacity, 6], Types.Real)
         self.subdom             = sim.add_static_array('subdom', [sim.ndims() * 2], Types.Real)
 
@@ -135,15 +134,20 @@ class BlockForest:
                 for r in For(self.sim, 0, self.nranks):
                     for aabb_id in For(self.sim, self.offsets[r], self.offsets[r] + self.naabbs[r]):
                         full_cond = None
+                        pbc_shifts = []
 
                         for d in range(self.sim.ndims()):
+                            cond_pbc_neg = position[i][d] - offset < self.sim.grid.min(d)
+                            cond_pbc_pos = position[i][d] + offset > self.sim.grid.max(d)
+                            d_pbc = Select(self.sim, cond_pbc_neg, -1, Select(self.sim, cond_pbc_pos, 1, 0))
+
+                            adj_pos = position[i][d] + d_pbc * self.sim.grid.length(d)
                             d_cond = ScalarOp.and_op(
-                                position[i][d] > self.aabbs[aabb_id][d * 2 + 0] + offset,
-                                position[i][d] < self.aabbs[aabb_id][d * 2 + 1] - offset)
+                                adj_pos > self.aabbs[aabb_id][d * 2 + 0] + offset,
+                                adj_pos < self.aabbs[aabb_id][d * 2 + 1] - offset)
 
-                            full_cond = d_cond if full_cond is None else \
-                                        ScalarOp.and_op(full_cond, d_cond)
+                            full_cond = d_cond if full_cond is None else ScalarOp.and_op(full_cond, d_cond)
+                            pbc_shifts.append(d_pbc)
 
                         for _ in Filter(self.sim, full_cond):
-                            pbc_shifts = [self.pbc[aabb_id][d] for d in range(self.sim.ndims())]
                             yield i, r, self.ranks[r], pbc_shifts