Clean files

MultiAgentsSim/simulation.py

@@ -6,9 +6,6 @@
 import matplotlib.pyplot as plt
 
 
-# import matplotlib
-# matplotlib.use('agg')
-
 # TODO : 
 @struct.dataclass
 class SimState:
@@ -27,9 +24,11 @@ def __init__(self, max_agents, grid_size):
         self.grid = self.init_grid(grid_size)
         self.max_agents = max_agents
 
+
     def init_grid(self, grid_size):
         return jnp.zeros((grid_size, grid_size), dtype=jnp.float32)
 
+
     def init_agents(self, num_agents, max_agents, key):
         if num_agents > max_agents:
             raise(ValueError("num_agents cannot exceed max_agents"))
@@ -44,14 +43,17 @@ def init_agents(self, num_agents, max_agents, key):
 
         return agents_pos, agents_states, num_agents
 
+
     def choose_random_action(self, agents_pos, key_a):
         return random.randint(key_a, agents_pos.shape, -1, 2)
 
+
     @partial(jit, static_argnums=(0, ))
     def move_agents(self, agents_pos, agents_movements):
         agents_pos += agents_movements
         return jnp.clip(agents_pos, 0, self.grid_size - 1)
 
+
     def add_agent(self, agents_pos, agents_states, num_agents, key):
         if num_agents < self.max_agents:
             agents_pos = agents_pos.at[num_agents].set(*random.randint(key, (1, 2), 0, self.grid_size))
@@ -64,6 +66,7 @@ def add_agent(self, agents_pos, agents_states, num_agents, key):
 
         return agents_pos, agents_states, num_agents
 
+
     def remove_agent(self, num_agents):
         if num_agents <= 0:
             print("There is no agents to remove")
@@ -72,29 +75,9 @@ def remove_agent(self, num_agents):
             print(f"Removed agent {num_agents + 1}")
         return num_agents
 
-    @staticmethod
-    def visualize_sim(grid, agents_pos, num_agents, color):
-        if not plt.fignum_exists(1):
-            plt.ion()
-            plt.figure(figsize=(10, 10))
-
-        plt.clf()
-
-        plt.imshow(grid, cmap="viridis", origin="upper")
-        plt.scatter(
-            agents_pos[:num_agents, 0], agents_pos[:num_agents, 1], color=color, marker="o", label="Agents"
-        )
-        plt.title("Multi-Agent Simulation")
-        plt.xlabel("X-axis")
-        plt.ylabel("Y-axis")
-        plt.legend()
-
-        plt.draw()
-        plt.pause(0.001)
-
 
     @staticmethod
-    def visualize_simzz(grid, agents_pos, num_agents, color):
+    def visualize_sim(grid, agents_pos, num_agents, color):
         if not plt.fignum_exists(1):
             plt.ion()
             plt.figure(figsize=(10, 10))

MultiAgentsSim/utils/network.py

@@ -0,0 +1,8 @@
+import pickle
+import socket 
+
+servers = ['10.204.2.229', '10.204.2.189', '192.168.1.24']
+server_idx = 0
+
+SERVER = servers[server_idx]
+

client.py

@@ -1,18 +1,14 @@
 import socket
 import pickle
-import threading
 
 from MultiAgentsSim.simulation import Simulation
+from MultiAgentsSim.utils.network import SERVER
 
-SERVER = '10.204.2.189'
-SERVER = '192.168.1.24'
 
 PORT = 5050
 ADDR = (SERVER, PORT)
-
 DATA_SIZE = 4096
 
-
 client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 client.connect(ADDR)
 print(f"Connected to {ADDR}")
@@ -23,7 +19,7 @@
 client.send(response.encode())
 print(f"responded: {response}")
 
-def receive():
+def receive_loop():
     while True:
         try:
             raw_data = client.recv(DATA_SIZE)
@@ -38,8 +34,4 @@ def receive():
             client.close()
             break
 
-# Matplotlib intreactive doesn't work outside the main thread
-# receive_thread = threading.Thread(target=receive)
-# receive_thread.start()
-
-receive()
+receive_loop()

server.py

@@ -7,14 +7,9 @@
 
 from MultiAgentsSim.simulation import Simulation
 from MultiAgentsSim.agents import Agents
+from MultiAgentsSim.utils.network import SERVER
 
 
-# Create and start the server 
-SERVER = '10.204.2.189'
-# SERVER = socket.gethostbyname(socket.gethostname())
-# print(f"{SERVER = }")
-
-SERVER = '192.168.1.24'
 PORT = 5050
 ADDR = (SERVER, PORT)
 DATA_SIZE = 4096 # size of data that is being transfered at each timestep
@@ -36,7 +31,6 @@
 
 
 key = jax.random.PRNGKey(SEED)
-
 sim = Simulation(MAX_AGENTS, GRID_SIZE)
 agents = Agents(MAX_AGENTS, GRID_SIZE)
 
@@ -65,10 +59,10 @@ def update_latest_data():
         timestep += 1
         time.sleep(STEP_DELAY)
 
+
 def update_color(new_color):
     global latest_data
     latest_data = latest_data[:5] + (new_color,) + (latest_data[6],)
-    print(f"entered the update color function and changed it")
 
 
 def handle_client(client, addr):
@@ -95,10 +89,8 @@ def handle_client(client, addr):
             while True:
                 try:
                     new_color = pickle.loads(client.recv(DATA_SIZE))
-                    print(f"{color = } before")
                     print(f"received new color {new_color} from client")
                     update_color(new_color)
-                    print(f"{color = } after")
                 except socket.error as e:
                     print(f"error: {e}")
                     client.close()

update_client.py

@@ -1,13 +1,12 @@
 import socket
 import pickle
 
-SERVER = '10.204.2.189'
-SERVER = '192.168.1.24'
+from MultiAgentsSim.utils.network import SERVER
+
 
 PORT = 5050
 DATA_SIZE = 4096
 
-
 update_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 update_client.connect((SERVER, PORT))
 print("Connected to server")
@@ -18,14 +17,17 @@
 update_client.send(response.encode())
 print(f"responded: {response}")
 
-while True:
-    try:
-        color = input("Enter a color in rgb format: ")
-        tuple_color = tuple(float(value) for value in color.split(" "))
-        print(f"{tuple_color = }")
-        update_client.send(pickle.dumps(tuple_color))
-        print(f"sent {tuple_color} to server")
-    except socket.error as e:
-        print(f"error: {e}")
-        update_client.close()
-        break
+def update_color_loop():
+    while True:
+        try:
+            color = input("Enter a color in rgb format: ")
+            tuple_color = tuple(float(value) for value in color.split(" "))
+            print(f"{tuple_color = }")
+            update_client.send(pickle.dumps(tuple_color))
+            print(f"sent {tuple_color} to server")
+        except socket.error as e:
+            print(f"error: {e}")
+            update_client.close()
+            break
+
+update_color_loop()