From 1ba57ff2d0019a83600ab476afd0aa79e98176b2 Mon Sep 17 00:00:00 2001
From: Rafael Ravedutti <rafaelravedutti@gmail.com>
Date: Tue, 6 Feb 2024 16:35:26 +0100
Subject: [PATCH] Fix PBC check

Signed-off-by: Rafael Ravedutti <rafaelravedutti@gmail.com>
---
 src/pairs/sim/domain_partitioning.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/pairs/sim/domain_partitioning.py b/src/pairs/sim/domain_partitioning.py
index 6adaa85..8a496f2 100644
--- a/src/pairs/sim/domain_partitioning.py
+++ b/src/pairs/sim/domain_partitioning.py
@@ -137,15 +137,18 @@ class BlockForest:
                         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)
+                            aabb_min = self.aabbs[aabb_id][d * 2 + 0]
+                            aabb_max = self.aabbs[aabb_id][d * 2 + 1]
+                            center = aabb_min + (aabb_max - aabb_min) * 0.5
+                            dist = position[i][d] - center
+                            d_length = self.sim.grid.length(d)
+
+                            cond_pbc_neg = dist >  (d_length * 0.5)
+                            cond_pbc_pos = dist < -(d_length * 0.5)
                             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(
-                                adj_pos > self.aabbs[aabb_id][d * 2 + 0] + offset,
-                                adj_pos < self.aabbs[aabb_id][d * 2 + 1] - offset)
-
+                            adj_pos = position[i][d] + d_pbc * d_length
+                            d_cond = ScalarOp.and_op(adj_pos > aabb_min - offset, adj_pos < aabb_max + offset)
                             full_cond = d_cond if full_cond is None else ScalarOp.and_op(full_cond, d_cond)
                             pbc_shifts.append(d_pbc)
 
-- 
GitLab