Periodic boundary conditions in 2D

3_3_periodic_boundary_conditions.py

import numpy as np

# Periodic boundary conditions 2D
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(a)
# 5: 2, 8, 4, 6 
i, j = 1, 1
print(f"{a[i, j]}'s nearest neighbor" )
print(a[i-1, j], a[i+1, j], a[i, j-1], a[i, j+1])

# 1: 7, 4, 2, 3
i, j = 0, 0
print(f"{a[i, j]}'s nearest neighbor" )
print(a[i-1, j], a[i+1, j], a[i, j+1], a[i, j-1])

# 9: 6, 3, 8, 7
i, j = 2, 2
ii = 0 if i + 1 == len(a[i]) else i
jj = 0 if j + 1 == len(a[:,j]) else j
print(f"{a[i, j]}'s nearest neighbor" )
print(a[i-1, j], a[ii, j], a[i, j-1], a[i, jj])
# np.roll
print(np.roll(a, 1, axis=0)[i, j], np.roll(a, -1, axis=0)[i, j], np.roll(a, 1, axis=1)[i, j], np.roll(a, -1, axis=1)[i, j])