How to get the neighbors for each agent?

[GuiyangLuo]

Hi, at first, I want to express my thanks to this awesome framework!

I want to obtain the neighbors of each agent, i.e., the spacial local view include the observed agents of both group 1 and group 2, I want to obtain the IDs and positions of all it's neighbors.

[merrymercy]

cc @KornbergFresnel

[outlace]

This should get you half way there. For a given group, this will return the index positions:

from scipy.spatial.distance import cityblock

def get_neighbors(j,pos_list,r=6):
    """
    Given [x,y] positions of all agents in `pos_list`,
    return indices of agents that are within the radius of agent j (i.e. the j'th index in `pos_list`)
    """
    neighbors = []
    pos_j = pos_list[j]
    for i,pos in enumerate(pos_list):
        if i == j:
            continue
        dist = cityblock(pos,pos_j)
        if dist < r:
            neighbors.append(i)
    return neighbors

>>> get_neighbors(5,env.get_pos(team1))
[27, 33, 34, 58]

where team1 is one of the group handles and env is an environment instance. The env.get_pos(...) method returns a list of [x,y] positions for each agent in that group.