pandas - Determine number of all the immediately adjacent points in python

Question

I have grid of point and want to calculate number of red and blue point surrounded by each point

for example

For point( X marked )

blue points = 6

red points = 2

I want to calculate it for the whole data set

import numpy as np
import pandas as pd    

x = np.linspace(1,50)
y = np.linspace(1,50)

GRID = np.meshgrid(x,y)
grid_colors = 1* ( np.random.random(GRID[0].size) > .8 )
sample_data = pd.DataFrame( {'X': GRID[0].flatten(), 'Y':GRID[1].flatten(), 'grid_color' : grid_colors})

sample_data.plot.scatter(x="X",y='Y', c='grid_color', colormap='bwr', figsize=(10,10))

red_points = sample_data[sample_data.grid_color == 1]
blue_points = sample_data[sample_data.grid_color != 1]

question from:https://stackoverflow.com/questions/66060235/determine-number-of-all-the-immediately-adjacent-points-in-python

Answer 1

scipy.signal.convolve2d seems to be perfect fit for this application.

color_matrix = grid_colors.reshape(50, 50)
from scipy import signal
adjacent_boxes = np.array([[1,1,1], [1,0,1], [1,1,1]])
red_dots = signal.convolve2d(color_matrix, adjacent_boxes, mode='same', boundary='fill')
blue_dots = signal.convolve2d(1 - color_matrix, adjacent_boxes, mode='same', boundary='fill')

red_dots
>>>
array([[1, 2, 2, ..., 0, 0, 0],
       [0, 2, 4, ..., 0, 0, 0],
       [1, 1, 2, ..., 1, 0, 0],
       ...,
       [1, 2, 1, ..., 1, 1, 0],
       [2, 3, 0, ..., 1, 1, 1],
       [1, 2, 0, ..., 1, 0, 0]])

blue_dots
>>>
array([[2, 3, 3, ..., 5, 5, 3],
       [5, 6, 4, ..., 8, 8, 5],
       [4, 7, 6, ..., 7, 8, 5],
       ...,
       [4, 6, 7, ..., 7, 7, 5],
       [3, 5, 8, ..., 7, 7, 4],
       [2, 3, 5, ..., 4, 5, 3]])

red_dots + blue_dots
>>>
array([[3, 5, 5, ..., 5, 5, 3],
       [5, 8, 8, ..., 8, 8, 5],
       [5, 8, 8, ..., 8, 8, 5],
       ...,
       [5, 8, 8, ..., 8, 8, 5],
       [5, 8, 8, ..., 8, 8, 5],
       [3, 5, 5, ..., 5, 5, 3]])

To have it inside dataframe, (thanks to @patrick artner) you can use:

sample_data['red_dots'] = red_dots.flatten()
sample_data['blue_dots'] = blue_dots.flatten()
print(sample_data)
         X     Y  grid_color  red_dots  blue_dots
0      1.0   1.0           0         1          2
1      2.0   1.0           0         2          3
2      3.0   1.0           1         2          3
...    ...   ...         ...       ...        ...
2498  49.0  50.0           0         0          5
2499  50.0  50.0           0         0          3

[2500 rows x 5 columns]