Random-Walk-Visualization / multi_agent_2D.py
DebasishDhal99's picture
Create multi_agent_2D.py
f115572
raw
history blame
4.89 kB
import matplotlib.pyplot as plt
import random
import numpy as np
import pandas as pd
from matplotlib.lines import Line2D
def single_random_walk(iters, agent_number, ax, step_size = 1, random_seed = None):
# random.seed(random_seed)
if random_seed:
random.seed(random_seed)
iters = int(iters)
directions = ['east', 'north', 'west', 'south']
start_point = [0, 0]
def distance_from_start(final_coord, start_coord, round_to=2):
return round(np.sqrt((final_coord[0] - start_coord[0])**2 + (final_coord[1] - start_coord[1])**2), round_to)
def step_addition(old_coord, step):
return [sum(x) for x in zip(old_coord, step)]
def step_determination():
direction = random.choice(directions)
if direction == 'east':
return [1*step_size, 0]
elif direction == 'west':
return [-1*step_size, 0]
elif direction == 'north':
return [0, 1*step_size]
elif direction == 'south':
return [0, -1*step_size]
coordinate_list = [start_point]
for _ in range(iters):
new_step = step_determination()
new_coordinate = step_addition(coordinate_list[-1], new_step)
coordinate_list.append(new_coordinate)
x = [i[0] for i in coordinate_list]
y = [i[1] for i in coordinate_list]
df = pd.DataFrame({'x':x,'y':y})
#Add the axis from the argument to the figure
base_marker_size = 10
markersize = base_marker_size / np.sqrt(iters)
plot = ax.plot(x, y, marker='o', markersize=markersize, linestyle='None', alpha=0.5, label = 'Agent {i}'.format(i=agent_number+1))
color = plot[0].get_color()
ax.plot(x[-1], y[-1], marker='o', markersize=5, color = 'black')
ax.text(x[-1], y[-1], 'End {i}'.format(i=agent_number+1), color = 'black', alpha=1.0)
return ax, df, color
def multi_agent_walk(agent_count, iters, step_size = 1, random_seed = None):
assert agent_count >= 1, "Number of agents must be >= than 1"
def displacement_calc(df):
x1,y1 = df.iloc[0]
x2,y2 = df.iloc[-1]
return np.round(np.sqrt((x2-x1)**2 + (y2-y1)**2),1)
if random_seed is None:
random_seed = random.randint(0,1000000)
assert type(random_seed) == int, "Random seed must be an integer"
#Generate a list of random seeds for each agent
random.seed(random_seed)
random_numbers = [random.randint(0,100000) for _ in range(agent_count)]
fig, ax = plt.subplots(figsize=(8,8))
color_list = []
for i in range(agent_count):
if i == 0:
ax, df, color = single_random_walk(iters=iters, ax=ax, step_size=step_size, agent_number=i, random_seed=random_numbers[i])
color_list.append(color)
else:
ax, df_new, color = single_random_walk(iters=iters, ax=ax, step_size=step_size, agent_number=i, random_seed=random_numbers[i])
df = pd.concat([df,df_new], axis=1)
x_columns = [f'x{i}' for i in range(1, i+2)]
y_columns = [f'y{i}' for i in range(1, i+2)]
new_column_names = [val for pair in zip(x_columns, y_columns) for val in pair]
df.columns = new_column_names
color_list.append(color)
ax.plot(0,0, marker='X', markersize=8, color='black')
ax.text(0, 0, 'Start (0,0)')
plt.grid()
plt.title('Random 2D Walk with {} agents\n #Steps = {}, Step size = {}, random seed = {}\nAll agents start from the origin'.format(agent_count, iters, step_size, random_seed))
displacement = [displacement_calc(df.iloc[:,[i,i+1]]) for i in range(0,agent_count*2,2)]
end_point = [(df.iloc[-1,i]) for i in range(0,agent_count*2,2)]
end_point = [(df.iloc[-1,i], df.iloc[-1,i+1]) for i in range(0,agent_count*2,2)]
agent_number = [i+1 for i in range(agent_count)]
legend_df = pd.DataFrame({'#':agent_number, 'dis.':displacement, 'End Point':end_point, })
info_box = legend_df.to_string(index=False)
ax.text(1.01, 0.99, info_box,
transform=ax.transAxes,
verticalalignment='top',
bbox=dict(boxstyle='round', facecolor='white', alpha=0.5)
)
lines = []
for i in range(len(color_list)):
lines.append(Line2D([0], [0], color=color_list[i], lw=9, linestyle=':'))
labels = [f'Agent {i+1}' for i in range(len(color_list))]
plt.legend(lines, labels,
loc='best',
handlelength=1.01,
handletextpad=0.21,
fancybox=True,
fontsize=10,
)
fig.canvas.draw()
image_array = np.array(fig.canvas.renderer.buffer_rgba())
try:
return image_array, df
except:
return image_array, None
# _, df = multi_agent_walk(agent_count=9, iters=1e5, step_size=1, random_seed=123);