diff --git a/src/pairs/sim/domain_partitioning.py b/src/pairs/sim/domain_partitioning.py index 654e98ceb7993b092113ea66e2d17318b4ce3b49..f3c6f4ea885c88ad3550512ae2aad21571b6cc47 100644 --- a/src/pairs/sim/domain_partitioning.py +++ b/src/pairs/sim/domain_partitioning.py @@ -11,7 +11,7 @@ from pairs.sim.lowerable import Lowerable class DimensionRanges: def __init__(self, sim): - self.sim = sim + self.sim = sim self.neighbor_ranks = sim.add_static_array('neighbor_ranks', [sim.ndims() * 2], Types.Int32) self.pbc = sim.add_static_array('pbc', [sim.ndims() * 2], Types.Int32) self.subdom = sim.add_static_array('subdom', [sim.ndims() * 2], Types.Real) @@ -23,20 +23,37 @@ class DimensionRanges: return [step * 2 + 0, step * 2 + 1] def ghost_particles(self, step, position, offset=0.0): - if self.sim._pbc[step] is False: - return - # Particles with one of the following flags are ignored flags_to_exclude = (Flags.Infinite | Flags.Global) - for i in For(self.sim, 0, self.sim.nlocal + self.sim.nghost): - for _ in Filter(self.sim, ScalarOp.cmp(self.sim.particle_flags[i] & flags_to_exclude, 0)): - j = step * 2 + 0 - for _ in Filter(self.sim, position[i][step] < self.subdom[j] + offset): - yield i, j, self.neighbor_ranks[j], [0 if d != step else self.pbc[j] for d in range(self.sim.ndims())] - - for i in For(self.sim, 0, self.sim.nlocal + self.sim.nghost): - for _ in Filter(self.sim, ScalarOp.cmp(self.sim.particle_flags[i] & flags_to_exclude, 0)): - j = step * 2 + 1 - for _ in Filter(self.sim, position[i][step] > self.subdom[j] - offset): - yield i, j, self.neighbor_ranks[j], [0 if d != step else self.pbc[j] for d in range(self.sim.ndims())] + def next_neighbor(self, j, step, position, offset, flags_to_exclude): + particle_flags = self.sim.particle_flags + for i in For(self.sim, 0, self.sim.nlocal + self.sim.nghost): + for _ in Filter(self.sim, ScalarOp.cmp(particle_flags[i] & flags_to_exclude, 0)): + for _ in Filter(self.sim, position[i][step] < self.subdom[j] + offset): + pbc_shifts = [0 if d != step else self.pbc[j] for d in range(self.sim.ndims())] + yield i, j, self.neighbor_ranks[j], pbc_shifts + + + + def prev_neighbor(self, j, step, position, offset, flags_to_exclude): + particle_flags = self.sim.particle_flags + j = step * 2 + 1 + for i in For(self.sim, 0, self.sim.nlocal + self.sim.nghost): + for _ in Filter(self.sim, ScalarOp.cmp(particle_flags[i] & flags_to_exclude, 0)): + for _ in Filter(self.sim, position[i][step] > self.subdom[j] - offset): + pbc_shifts = [0 if d != step else self.pbc[j] for d in range(self.sim.ndims())] + yield i, j, self.neighbor_ranks[j], pbc_shifts + + if self.sim._pbc[step]: + yield from next_neighbor(self, step * 2 + 0, step, position, offset, flags_to_exclude) + yield from prev_neighbor(self, step * 2 + 1, step, position, offset, flags_to_exclude) + + else: + j = step * 2 + 0 + for _ in Filter(self.sim, ScalarOp.inline(ScalarOp.cmp(self.pbc[j], 0))): + yield from next_neighbor(self, j, step, position, offset, flags_to_exclude) + + j = step * 2 + 1 + for _ in Filter(self.sim, ScalarOp.inline(ScalarOp.cmp(self.pbc[j], 0))): + yield from prev_neighbor(self, j, step, position, offset, flags_to_exclude)