Create multi_agent_2D.py

multi_agent_2D.py

@@ -0,0 +1,136 @@
+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);
+
+