The lines periodicity is not working like the columns periodicity. I think you should first get the 2 cells on each side and then move up and down. I have tried this and it seems to work :
def celdas_vecinas(cell, n) : last_row = n * (cell // n) left_cell = last_row + ( cell - last_row - 1 ) % n right_cell = last_row + ( cell - last_row + 1 ) % n line = np.array( [ left_cell, cell, right_cell ] ) return np.mod( [ line + n, line, line - n ], n**2)(I removed my previous answer as I messed up in indeces)